#!/usr/bin/env python3 import argparse import glob import os import subprocess import sys import threading from contextlib import contextmanager CONFIG = {} DC_ENV = {} def read_proc_callback(proc, output_handler=lambda line: line): for line in iter(proc.stdout.readline, b''): output_handler(line) def read_config(fpath='../.env.dev'): with open(fpath, 'r') as f: content = f.read() config = {} for line in content.splitlines(): line = line.strip() if line.startswith('#'): continue key, val = line.split('=', 1) config[key] = val return config def is_container_running(container): return bool(subprocess.check_output(['docker', 'ps', '-q', '--filter', f'name={container}'], env=DC_ENV)) def get_containers_ids(*additional_docker_ps_args): names = subprocess.check_output(['docker', 'ps', '-q', *additional_docker_ps_args], env=DC_ENV) return names.decode('utf-8').split() def stop(containers=None): """ Stop passed containers or all containers if no names passed :param containers: list|tuple names or string name or None, meaning all :return: list of stopped containers names""" if containers is None: containers = get_containers_ids() if not isinstance(containers, (list, tuple)): containers = [containers] if containers: arguments = ['docker', 'stop', *containers] print(f'Stopping containers {containers}') proc = subprocess.run(arguments, env=DC_ENV) return containers def down(): """ Stop and remove containers """ arguments = ['docker-compose', '-f', 'compose/django.yml', '-f', 'compose/selenium.yml', '-f', 'compose/flower.yml', '-f', 'compose/pgadmin.yml', '-f', 'compose/el.yml', '-f', 'compose/correlator.yml', '--env-file=../.env.dev', '-f', 'compose/kibana.yml', '-f', 'compose/license.yml', 'down'] print('Docker-compose down') subprocess.run(arguments, env=DC_ENV) def exec_django_container_cmd(*cmds): global CONFIG subprocess.run(['docker', 'exec', '-ti', CONFIG.get('DOCKER_DJANGO_NAME'), *cmds], env=DC_ENV) def wait_process_stop(proc, timeout=None): """ Wait for process which supports interruption, i.e. it can stop itself on interruption """ try: try: return proc.wait(timeout=timeout) except KeyboardInterrupt: print('Stopping process gracefully, repeat interruption to stop immediately') proc.wait() except: proc.kill() proc.wait() raise def wait_till_output_bytes_line(proc, bytes_lines_to_wait=[b'Starting Gnuicorn server', b'Starting Django dev server'], print_output=True, reraise=False): if not isinstance(bytes_lines_to_wait, (tuple, list)): bytes_lines_to_wait = [bytes_lines_to_wait] try: print('Waiting for docker ready') for line in iter(proc.stdout.readline, b''): if print_output: print(line.decode('utf-8')[:-1]) if any((l in line for l in bytes_lines_to_wait)): return True return False except KeyboardInterrupt: print('Stopping process gracefully, repeat interruption to stop immediately') return False @contextmanager def docker_on(use_gunicorn=False, use_selenium=False, use_el=False, use_correlator=False, use_kibana=False, use_flower=False, use_pgadmin=False, wait_till_loaded=True): global CONFIG arguments = ['docker-compose', '-f', 'compose/django.yml', '-f', 'compose/license.yml'] was_running = [] needed_conts = ('DOCKER_DJANGO_NAME',) was_running.extend([is_container_running(CONFIG.get(cont, 'None')) for cont in needed_conts]) if use_selenium: needed_conts = ('DOCKER_FIREFOX_NODE_NAME', 'DOCKER_CHROME_NODE_NAME', 'DOCKER_SELENIUM_NAME') was_running.extend([is_container_running(CONFIG.get(cont, 'None')) for cont in needed_conts]) arguments.append('-f') arguments.append('compose/selenium.yml') if use_flower: needed_conts = ('DOCKER_FLOWER_NAME',) was_running.extend([is_container_running(CONFIG.get(cont, 'None')) for cont in needed_conts]) arguments.append('-f') arguments.append('compose/flower.yml') if use_pgadmin: needed_conts = ('DOCKER_PGADMIN_SERVER',) was_running.extend([is_container_running(CONFIG.get(cont, 'None')) for cont in needed_conts]) arguments.append('-f') arguments.append('compose/pgadmin.yml') if use_el: needed_conts = ('DOCKER_ELASTIC_NAME', 'DOCKER_VECTOR_NAME') was_running.extend([is_container_running(CONFIG.get(cont, 'None')) for cont in needed_conts]) arguments.append('-f') arguments.append('compose/el.yml') arguments.append('--env-file=../.env.dev') if use_correlator: arguments.append('-f') arguments.append('compose/correlator.yml') if use_kibana: needed_conts = ('DOCKER_KIBANA_NAME',) was_running.extend([is_container_running(CONFIG.get(cont, 'None')) for cont in needed_conts]) arguments.append('-f') arguments.append('compose/kibana.yml') if not use_gunicorn: DC_ENV['SERVE'] = 'django' was_running = all([was_running, *was_running]) arguments.append('up') proc = None if not was_running: print(f'Starting containers') kwargs = {} if wait_till_loaded: kwargs = dict(stdout=subprocess.PIPE, stderr=subprocess.STDOUT) proc = subprocess.Popen(arguments, env=DC_ENV, **kwargs) if wait_till_loaded: if not wait_till_output_bytes_line(proc): stop() raise KeyboardInterrupt try: if proc and wait_till_loaded: # prevent process queue overflow def drop_lines(line): pass threading.Thread(target=read_proc_callback, args=(proc, drop_lines)).start() yield proc except: if proc: proc.kill() proc.wait() raise finally: if not was_running: stop() if proc: wait_process_stop(proc) def test(enable_warning=False, add_args=[]): """ Call process as subprocess.call, but on interruption wait for termination of process """ global CONFIG test_args = ['docker-compose', '-f', 'compose/django.yml', '-f', 'compose/license.yml', '-f', 'compose/selenium.yml', 'exec', CONFIG.get('DOCKER_DJANGO_NAME'), 'pytest'] if not enable_warning: test_args.append('--disable-warnings') test_args.extend(add_args) with docker_on(use_selenium=True, wait_till_loaded=True): print('Starting tests') test_proc = subprocess.Popen( test_args, env=DC_ENV, stdout=subprocess.STDOUT, stderr=subprocess.STDOUT) wait_process_stop(test_proc) sys.exit(test_proc.returncode) def ci_test(add_args): """ Used to run tests in CI """ global CONFIG test_args = ['docker-compose', '-f', 'compose/django.yml', '-f', 'compose/license.yml', '-f', 'compose/selenium.yml', 'exec', '-T', CONFIG.get('DOCKER_DJANGO_NAME'), 'pytest'] test_args.extend(add_args) with docker_on(use_selenium=True, wait_till_loaded=True): print('Starting tests') test_proc = subprocess.Popen(test_args, env=DC_ENV) wait_process_stop(test_proc) print("Exit code:", test_proc.returncode) sys.exit(test_proc.returncode) def ci_test_coverage(add_args): """ Used to check test coverage""" global CONFIG DC_ENV["COVERAGE_FILE"] = "public/.coverage" test_args = ['docker-compose', '-f', 'compose/django.yml', '-f', 'compose/license.yml', '-f', 'compose/selenium.yml', 'exec', '-T', CONFIG.get('DOCKER_DJANGO_NAME'), 'coverage', "run", "-m", "pytest", "-m", "'merge or selenium'"] report_args = ['docker-compose', '-f', 'compose/django.yml', '-f', 'compose/license.yml', 'exec', '-T', CONFIG.get('DOCKER_DJANGO_NAME'), 'coverage', "html", "-d", "public/test_coverage/"] test_args.extend(add_args) result = 0 with docker_on(use_selenium=True, wait_till_loaded=True): print(f"Starting tests {test_args}") test_proc = subprocess.Popen(test_args, env=DC_ENV) wait_process_stop(test_proc) result = test_proc.returncode with docker_on(use_selenium=True, wait_till_loaded=True): print(f"Processing report {report_args}") test_report = subprocess.Popen(report_args, env=DC_ENV) wait_process_stop(test_report) sys.exit(result) def parse_args(argv=sys.argv[1:]): parser = argparse.ArgumentParser( description=f'Project starter. To open web ui http://localhost:{CONFIG.get("WEB_UI_PORT", 9090)}/ with admin:nimda') sub_parsers = parser.add_subparsers(title='Commands', dest='cmd', required=True) run_parser = sub_parsers.add_parser(name='run', help='Run custom configuration') run_parser.add_argument('--with-selenium', dest='use_selenium', action='store_true', help='Add selenium') run_parser.add_argument('--with-flower', dest='use_flower', action='store_true', help='Add flower. Use http://localhost:5555') run_parser.add_argument('--with-pgadmin', dest='use_pgadmin', action='store_true', help='Add Pgadmin. Use http://localhost:5050 with pgadmin4@pgadmin.org : admin') run_parser.add_argument('--with-ev', dest='use_ev', action='store_true', help='Add Vector and Elasticsearch') run_parser.add_argument('--with-cev', dest='use_cev', action='store_true', help='Add Vector, Elasticsearch and correlator') run_parser.add_argument('--with-kibana', dest='use_kibana', action='store_true', help='Add Kibana. Use http://localhost:5601 with elastic:changeme') run_parser.add_argument('--with-gunicorn', dest='use_gunicorn', action='store_true', help='Use gunicorn server instead of django') run_parser.add_argument('--with-evk', dest='use_evk', action='store_true', help='Add Vector, Elasticsearch and Kibana') run_parser.add_argument('--with-cevk', dest='use_cevk', action='store_true', help='Add Vector, Elasticsearch, correlator and Kibana') test_parser = sub_parsers.add_parser(name='test', aliases=['tests'], help='Run tests') test_parser.add_argument('--enable-warning', dest='enable_warning', action='store_true', help='Enable warning for tests') ci_test_parser = sub_parsers.add_parser(name="ci_test", help="Run tests in CI. Differs from the test function in the way of running docker compose. Use for CI only.") ci_test_coverage = sub_parsers.add_parser(name="ci_test_coverage", help="Run test coverage in CI. Differs from the test function in the way of running docker compose. Use for CI only.") build_parser = sub_parsers.add_parser(name='build', aliases=['rebuild'], help='Rebuild containers') build_parser.add_argument('--no-cache', dest='no_cache', action='store_true', help='Rebuild without cache') stop_parser = sub_parsers.add_parser(name='stop', help='Stop containers') down_parser = sub_parsers.add_parser(name='down', help='Stop containers and remove containers') django_restart_db_parser = sub_parsers.add_parser(name='restart-django', aliases=['reload-django'], help='Restart django container') sh_parser = sub_parsers.add_parser(name='sh', help='Run shell in django container') psh_parser = sub_parsers.add_parser(name='psh', help='Run python django shell in django container') psql_parser = sub_parsers.add_parser(name='psql', help='Run postgresql shell') clear_db_parser = sub_parsers.add_parser(name='clear-db', aliases=['db-clear'], help="Remove migrations which are not in git and drop container's data and tables from database") clear_db_parser.add_argument('-n', '--dry-run', dest='dry_run', action='store_true', help='Show actions without actual deletion of info') webui_parser = sub_parsers.add_parser(name='ui', help='Open web UI') webpdb_parser = sub_parsers.add_parser(name='wpdb', help='Open web PDB debugger') args, unknownargs = parser.parse_known_args(argv) return args, unknownargs def main(): global CONFIG, DC_ENV CONFIG.update(read_config()) DC_ENV = {**os.environ.copy(), **CONFIG} args, unknownargs = parse_args() if args.cmd == 'run': if args.use_cevk: args.use_evk, args.user_cev = True, True if args.use_cev: args.use_ev = True if args.use_evk: args.use_ev, args.use_kibana = True, True with docker_on(args.use_gunicorn, args.use_selenium, args.use_ev, args.use_cevk, args.use_kibana, args.use_flower, args.use_pgadmin, wait_till_loaded=False) as proc: if proc: wait_process_stop(proc) elif args.cmd in ('test', 'tests'): test(args.enable_warning, unknownargs) elif args.cmd in ('ci_test'): ci_test(unknownargs) elif args.cmd in ("ci_test_coverage"): ci_test_coverage(unknownargs) elif args.cmd in ('build', 'rebuild'): stop() arguments = ['docker-compose', '-f', 'compose/django.yml', '-f', 'compose/license.yml', '-f', 'compose/selenium.yml', '-f', 'compose/pgadmin.yml', '--env-file=../.env.dev', '-f', 'compose/el.yml', '-f', 'compose/correlator.yml', '-f', 'compose/kibana.yml', 'build'] if args.no_cache: arguments.append('--no-cache') proc = subprocess.run(arguments, env=DC_ENV) elif args.cmd == 'stop': stop() elif args.cmd == 'down': down() elif args.cmd in ('restart-django', 'reload-django'): subprocess.run(['docker', 'restart', CONFIG.get('DOCKER_DJANGO_NAME')], env=DC_ENV) elif args.cmd in ('sh',): with docker_on(wait_till_loaded=True) as proc: exec_django_container_cmd('bash', *unknownargs) elif args.cmd in ('psh',): with docker_on(wait_till_loaded=True) as proc: exec_django_container_cmd(*['bash', '-c', 'python3 manage.py shell'], *unknownargs) elif args.cmd in ('psql',): with docker_on(wait_till_loaded=True) as proc: arguments = ['psql', '-U', CONFIG.get('POSTGRES_USER'), '-d', CONFIG.get("POSTGRES_DB")] subprocess.run(['docker', 'exec', '-ti', CONFIG.get('DOCKER_DB_NAME'), *arguments], env=DC_ENV) elif args.cmd in ('clear-db', 'db-clear'): print('Clearing migrations') git_clean_opts = '-xf' + (args.dry_run and 'n' or '') subprocess.run(['sudo', 'git', 'clean', git_clean_opts, *glob.glob("../*/migrations/")], env=DC_ENV) subprocess.run(['sudo', 'git', 'clean', git_clean_opts, 'compose/config/elk/vector/pipeline/armaif_*.toml'], env=DC_ENV) subprocess.run(['sudo', 'git', 'clean', git_clean_opts, 'compose/config/elk/vector/pipeline/endpoint_*.toml'], env=DC_ENV) print("Applying docker-compose down to all containers") if not args.dry_run: down() elif args.cmd == 'ui': import webbrowser webbrowser.open(f'http://localhost:{CONFIG.get("WEB_UI_PORT")}') elif args.cmd == 'wpdb': import webbrowser webbrowser.open(f'http://localhost:{CONFIG.get("WEB_PDB_PORT")}') if __name__ == '__main__': main()