__init__.py

import functools
import io
import os
import re
import shutil
import signal
import sys
import time
import getpass
from multiprocessing.pool import ThreadPool
from zipfile import ZipFile

from click.exceptions import Exit
from rich import print
from rich.markup import escape
import typer
import base64
import tarfile

from typing import List

from rich.prompt import Confirm

from player_cli.exploit import target
from .execution import print_exploit_execution, EXECUTION_STATUS_IS_FINAL, EXECUTION_STATUS_COLOR
from .exploit import get_all_histories, resolve_history, deactivate_history, activate_exploit, resolve_exploit, \
    print_history, print_logs, ResolveStrategy
from player_cli.util import request, WARN_STR, ERROR_STR, make_executable, parse_dockerfile_cmd, dt_to_local_str, \
    dt_from_iso, highlight_flags, blueify, magentify, NOTICE_STR
from player_cli.ctfconfig_wrapper import RUNLOCAL_TARGETS, ROUND_TIME
from .job import run_local_job
from .target import get_targets
from ..flags import poll_and_show_flags

app = typer.Typer(no_args_is_help=True)

app.add_typer(target.app, name='target', help='Manage exploit targets.')


@app.command('ls', help='List all exploits.')
def exploit_ls(
        history_ids: List[str] = typer.Argument(None, help='History ID(s).')
):
    histories = resolve_history(history_ids) if history_ids is not None and len(history_ids) > 0 else get_all_histories()

    if len(histories) == 0:
        print("No exploits found")
        return

    for history in histories:
        print_history(history)


@app.command('activate', help='Activate an exploit.', no_args_is_help=True)
def exploit_activate(
        exploit_id: str = typer.Argument(..., help=
        'Exploit ID or history ID. '
        'If an exploit ID is specified, activates that exploit. '
        'If a history ID is specified, activates the most recent exploit in '
        'the history if none is already activated.')
):
    exploit = resolve_exploit(exploit_id)
    return activate_exploit(exploit)


@app.command('deactivate', help='Deactivate an exploit.', no_args_is_help=True)
def exploit_deactivate(
        exploit_id: str = typer.Argument(..., help=
        'Exploit ID or history ID. '
        'If an exploit ID is specified, deactivates that exploit. '
        'If a history ID is specified, deactivates all exploits in the history.')
):
    history = resolve_history(exploit_id)
    deactivate_history(history)


@app.command('switch', help=
'Activate an exploit and deactivate all others in the history.', no_args_is_help=True)
def exploit_switch(
        exploit_id: str = typer.Argument(..., help='Exploit ID.')
):
    exploit = resolve_exploit(exploit_id)
    history = exploit['history']
    if history['id'] == exploit_id:
        print(f'{ERROR_STR}: you specified an exploit history id, not a specific exploit id.')
        print(f'{ERROR_STR}: Please select one of the following exploits to switch to.')
        print_history(history)
        raise typer.Exit(code=1)

    if exploit['active']:
        print(f'{WARN_STR}: exploit "{exploit_id}" is already active, doing nothing...')
        return

    deactivate_history(history)
    activate_exploit(exploit)

@app.command('create', help='Create an exploit history.', no_args_is_help=True)
def exploit_create(
        history_id: str = typer.Argument(..., help='History ID (the "friendly" exploit name).'),
        service: str = typer.Argument(..., help='The target service.')
):
    targets = get_targets(None)
    services = set([target['service'] for target in targets])

    if service not in services:
        print(f'{ERROR_STR}: unknown service "{service}". Available services: {magentify(", ".join(services))}.')
        raise typer.Exit(code=1)

    request('POST', 'exploit_history', data={
        'history_id': history_id,
        'service': service,
    })

    print(f"Exploit history {history_id} created!")


@app.command('logs', help='Show remote exploit logs.', no_args_is_help=True)
def exploit_logs(
        cmd_ids: List[str] = typer.Argument(..., metavar='EXPLOIT_ID...', help=
        'Exploit ID or history ID. '
        'If an exploit ID is specified, shows logs for that exploit. '
        'If a history ID is specified, shows logs for the active exploits in the history. '
        'You can specify multiple IDs and mix exploit and history IDs.'),
        limit: int = typer.Option(1, '-n', '--num', metavar='NUM', help=
        'Show logs for the last NUM ticks.')
):
    exploits = resolve_exploit(cmd_ids, ResolveStrategy.ACTIVE)

    print_logs(exploits, limit)


@app.command('upload', help='Upload an exploit.', no_args_is_help=True)
def exploit_upload(
        history_id: str = typer.Argument(..., help='History ID (the "friendly" exploit name).'),
        author: str = typer.Argument(..., help='The author.'),
        context: str = typer.Argument(..., help=
        'Path to a directory or tarball containing the Docker context. '
        'The Dockerfile must be top-level. '
        'Supported tarball compression formats: xz, bzip2, gzip, identity (no compression).'),
        y: bool = typer.Option(False, help='Switch to new exploit after uploading.'),
):
    if os.path.isdir(context):
        with io.BytesIO() as bio:
            with tarfile.open(fileobj=bio, mode='w:gz') as tar:
                tar.add(context, arcname='')
            bio.seek(0)
            context_data = bio.read()
    else:
        with open(context, 'rb') as f:
            context_data = f.read()

    exploit = request('POST', 'exploit', data={
        'history_id': history_id,
        'author': author,
        'context': base64.b64encode(context_data).decode(),
    })

    exploit_id = exploit['id']

    targets = get_targets(exploit['history']['service'], all_targets=False)

    job = request('POST', 'job', data={
        'targets': [target['id'] for target in targets],
        'exploit_id': exploit_id,
        'manual_id': None,
        'timeout': ROUND_TIME,
    })

    job_id = job['id']
    print(f"{exploit_id}: Waiting for initial job {job_id} to finish..", end='')
    job_status = None

    for i in range(120):
        job = request('GET', f'job/{job_id}')
        executions = job['executions']
        print(f".", end='')
        if job_status != job['status']:
            print(f" {EXECUTION_STATUS_COLOR[job['status']](job['status'])}", end='')
            if job['status'] == 'queued':
                print(f" (probably building right now)", end='')
        job_status = job['status']

        if job['status'] in EXECUTION_STATUS_IS_FINAL:
            job['timestamp'] = dt_from_iso(job['timestamp'])
            print("\n")
            for e in executions:
                print_exploit_execution(job, e)

            poll_and_show_flags([e['id'] for e in executions], timeout=2)
            break

        time.sleep(1)

    if job_status == 'finished':
        if not y:
            msg = f"{NOTICE_STR}: Do you want to activate the newly uploaded exploit {exploit_id}?"
            history = resolve_history(exploit_id)
            if len(previous_versions := [exploit['id'] for exploit in history['exploits'] if exploit['active']]) > 0:
                msg += f"\n{NOTICE_STR}: this would deactivate the previous version: {','.join(previous_versions)}"
            should_switch = Confirm.ask(msg)
        else:
            should_switch = True

        if should_switch:
            if len([exploit['id'] for exploit in history['exploits'] if exploit['active']]) > 0:
                exploit_switch(exploit_id)
            else:
                exploit_activate(exploit_id)
        else:
            print(f"{WARN_STR}: Did not activate exploit {exploit_id}.")

    elif job_status == 'failed':
        print(f"{ERROR_STR}: Job execution failed, not switching exploit automatically.")
        print(f"Please fix the issue and re-run the upload")
    else:
        print(f"{WARN_STR}: Job did not finish within 120 seconds, not switching exploit over automatically.")
        print(f"Check results manually with `{sys.argv[0]} exploit logs {exploit_id}`")
        print(f"Afterwards use `{sys.argv[0]} exploit switch {exploit_id}` to switch over central execution")


@app.command('download', help='Download an exploit.', no_args_is_help=True)
def exploit_download(
        exploit_id: str = typer.Argument(..., help='Exploit ID.'),
        path: str = typer.Argument(..., help='Output directory (will be created).'),
        overwrite: bool = typer.Option(False, '--overwrite', help=
        'Proceed even if the destination directory already exists.'),
        unsafe: bool = typer.Option(False, '--unsafe', help=
        'DANGEROUS: do not perform safety checks before extracting the downloaded archive.')
):
    resp = request('GET', f'exploit/{exploit_id}/download')

    data = base64.b64decode(resp['data'])

    with io.BytesIO(data) as bio:
        with tarfile.open(fileobj=bio) as tar:
            if not unsafe:
                for info in tar:
                    if info.name.startswith('/'):
                        print(f'{ERROR_STR}: unsafe archive: absolute path detected')
                        raise typer.Exit(code=1)
                    elif '/../' in f'/{info.name}/':
                        print(f'{ERROR_STR}: unsafe archive: path traversal detected')
                        raise typer.Exit(code=1)
                    elif info.islnk() or info.issym():
                        print(f'{ERROR_STR}: unsafe archive: link detected')
                        raise typer.Exit(code=1)

            # Check if the directory already exists and ask the user if 
            # overwriting is fine if overwrite is not explicitly specified
            if os.path.exists(path) and not overwrite:
                confirm_str = f'{ERROR_STR}: directory "{path}" already exists, do you want to overwrite it?'

                overwrite_answer = Confirm.ask(confirm_str)
                if not overwrite_answer:
                    raise typer.Exit(code=1)

            # Recursivly create directories, not raising FileExistsError on
            # existing directories(overwrite-case)
            os.makedirs(path, exist_ok=True)
            tar.extractall(path)


self_as_zip_path = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))


@app.command("template", help="Generate an exploit stub from a template.", no_args_is_help=True)
def exploit_template(
        template: str = typer.Argument(
            ...,
            help="Template to use, one of: python, ubuntu. "
                 "Optionally, you can specify a Docker tag (e.g., python:3.9-slim). "
                 "If none is specified, the latest will be used.",
        ),
        path: str = typer.Argument(..., help="Destination directory."),
        overwrite: bool = typer.Option(
            False,
            "--overwrite",
            help="Proceed even if the destination directory already exists.",
        ),
):
    if template.count(":") > 1:
        raise typer.BadParameter("Can't give more than one tag", param_hint="TEMPLATE")

    with ZipFile(self_as_zip_path, "r") as archive:
        templates = [info.filename.split("/")[1] for info in archive.infolist()
                     if info.is_dir() and info.filename.count("/") == 2 and info.filename.startswith("templates/")]
        template, tag = template.split(":") if ":" in template else (template, "latest")

        if template not in templates:
            raise typer.BadParameter(f'unknown template {template}', param_hint="TEMPLATE")

        for info in [x for x in archive.infolist() if x.filename.startswith(f"templates/{template}/")]:
            out_filename = os.path.join(path, info.filename[len(f"templates/{template}/"):])
            if info.is_dir():
                try:
                    os.mkdir(out_filename)
                except FileExistsError:
                    if overwrite:
                        print(
                            f'{WARN_STR}: directory "{out_filename}" already exists (proceeding anyway)'
                        )
                    else:
                        print(
                            f'{ERROR_STR}: directory "{out_filename}" already exists (use --overwrite to proceed anyway)'
                        )
                        raise typer.Exit(code=1)
            else:
                with open(out_filename, "wb") as outfile:
                    with archive.open(info.filename) as infile:
                        contents = infile.read()

                        if info.filename == f"templates/{template}/Dockerfile":
                            contents = re.sub(b"^FROM (.*):.*$", f"FROM \\1:{tag}".encode(), contents,
                                              flags=re.MULTILINE)

                        outfile.write(contents)
                # if file is executable
                if (info.external_attr >> 16) & 0o111 > 0:
                    make_executable(out_filename)

    print(f'Created exploit "{path}/" from template {template}:{tag}')


@app.command('runlocal', help='Run an exploit locally.', no_args_is_help=True)
def exploit_runlocal(
        path: str = typer.Argument(..., help=
        'Path to exploit executable, or to a directory containing a Dockerfile. '
        'In the latter case, will try to run the container command locally. '
        'The working directory will be the one where the Dockerfile is located.'),
        service: str = typer.Argument(..., help='The target service.'),
        target_ips: List[str] = typer.Option(RUNLOCAL_TARGETS, '--target', '-T', help=
        'Target to attack (you can specify this option multiple times).'),
        no_target_ips: List[str] = typer.Option([], '-N', '--no-target', help=
        'Target to not attack (you can specify this option multiple times).'),
        all_targets: bool = typer.Option(False, '--all-targets', help=
        'Attack all targets (overrides --target).'),
        ignore_exclusions: bool = typer.Option(False, help=
        'Ignore static exclusions, i.e. excluding our own vulnbox.'),
        timeout: int = typer.Option(30, '--timeout', '-t', help=
        'Timeout for a single exploit execution, in seconds.'),
        jobs: int = typer.Option(0, '--jobs', '-j', help=
        'Number of parallel jobs (0 for CPU count).'),
        limit: int = typer.Option(-1, '--limit', '-l', help=
        'Limit stdout printing to <limit> chars. '
        'Set to -1 if you want to see the whole output.'),
        count: int = typer.Option(0, '-c', '--count', help=
        'Number of attack rounds to perform (0 for infinite).')
):
    if os.path.isdir(path):
        try:
            with open(f'{path}/Dockerfile', 'r') as f:
                dockerfile = f.read()
        except FileNotFoundError:
            print(f'{ERROR_STR}: directory specified, but no Dockerfile found')
            raise typer.Exit(code=1)
        except IOError:
            print(f'{ERROR_STR}: error reading Dockerfile')
            raise typer.Exit(code=1)
        exe_args = parse_dockerfile_cmd(dockerfile)
        if exe_args is None:
            print(f'{ERROR_STR}: could not extract command from Dockerfile')
            raise typer.Exit(code=1)
        exe = shutil.which(exe_args[0])
        if exe is None:
            print(f'{ERROR_STR}: could not find executable for "{exe_args[0]}"')
            raise typer.Exit(code=1)
        workdir = path
    else:
        if not os.access(path, os.X_OK):
            print(f'{ERROR_STR}: exploit file is not executable')
            raise typer.Exit(code=1)
        exe = path
        exe_args = [path]
        workdir = '.'

    if jobs == 0:
        jobs = None

    manual_id = f"{os.uname().nodename}-{getpass.getuser()}-{os.path.basename(os.path.realpath(workdir))}"

    targets = get_targets(None)
    services = set([target['service'] for target in targets])

    if service not in services:
        print(f'{ERROR_STR}: unknown service "{service}". Available services: {magentify(", ".join(services))}.')
        raise typer.Exit(code=1)

    original_sigint_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
    pool = ThreadPool(jobs)
    signal.signal(signal.SIGINT, original_sigint_handler)

    job_func = functools.partial(run_local_job, exe=exe, args=exe_args, workdir=workdir, timeout=timeout)

    num_rounds = 0
    while True:
        job = None
        try:
            print(f'\\[{manual_id}] Attack round started')

            start = time.time()

            targets = get_targets(None, all_targets=all_targets, target_ips=target_ips,
                                  no_target_ips=no_target_ips, ignore_exclusions=ignore_exclusions)
            services = set([target['service'] for target in targets])
            targets = {target['id']: target for target in targets if target['service'] == service}
            if service not in services:
                print(f'{ERROR_STR}: No such service {service}, aborting...')
                raise typer.Exit(1)

            targets_summary = ', '.join(t['ip'] for t in targets.values())
            print(f'\\[{manual_id}] Attacking {len(targets)} targets: {targets_summary}')

            job = request('POST', 'job', data={
                'targets': list(targets.keys()),
                'exploit_id': None,
                'manual_id': manual_id,
                'timeout': ROUND_TIME,
            })

            for execution in job['executions']:
                execution['target'] = targets[execution['target_id']]
                execution['finished'] = False

            for result in pool.imap_unordered(job_func, job['executions']):
                host = result['target']['ip']
                extra = result['target']['extra']
                msg = result['msg']
                flags = result['flags']
                print(f'\\[{manual_id}] Execution completed on target {host}: {msg}')
                print(f'    Extra: {extra}')
                #print(f'    Submitted {len(flags)} flags:')
                #for flag in flags:
                #    dt = dt_to_local_str(dt_from_iso(flag["timestamp"]))
                #    print(f'        {dt} | ID {flag["id"]}: {flag["flag"]} ({flag["status"]})')

                if limit == -1:
                    print(f'    Execution output:')
                    # Assumes author wants to see the stdout in a beautiful way, so no repr
                    print(highlight_flags(escape(result['stdout']), blueify))
                elif result["stdout"] != '':
                    print(f'    Execution output: {magentify(escape(repr(result["stdout"][:limit])))}')
                if result["stderr"] != '':
                    print('\n'.join([f'[bold red]ERR[/bold red] {x}' for x in highlight_flags(escape(result['stderr']), blueify).split('\n')]))

            executions = job['executions']
            request('POST', f'job/{job["id"]}/finish')
            job = None

            poll_and_show_flags([execution['id'] for execution in executions])

            num_rounds += 1
            if count > 0 and num_rounds == count:
                break

            end = time.time()
            if end - start < ROUND_TIME:
                to_wait = ROUND_TIME - (end - start)
                print(f'\n\\[{manual_id}] Waiting {to_wait:.2f} seconds for next round')
                time.sleep(to_wait)
        except KeyboardInterrupt:
            print(f'\\[{manual_id}] Terminating...')
            pool.terminate()

            if job is not None:
                print(f"\\[{manual_id}] Cancelling job...")
                for execution in job['executions']:
                    if not execution['finished']:
                        request('POST', f'job/execution/{execution["id"]}/finish', data={
                            "stdout": '',
                            "stderr": '',
                            "status": 'cancelled'
                        })

                request('POST', f'job/{job["id"]}/finish', params={"status": 'cancelled'})
                print(f"\\[{manual_id}] Done")
            return