tester

#!/usr/bin/env python3

FLAG_RE = r"OOO{[^}]*}"
CTF_TAG = "dc2019f"
IMAGE_FMT = "%s"


import concurrent.futures
import contextlib
import subprocess
import argparse
import tempfile
import logging
import pathlib
import shutil
import random
import string
import json
import yaml
import time
import sys
import re
import os

logging.basicConfig()
_LOG = logging.getLogger("OOO")
_LOG.setLevel("DEBUG")

try:
    import ooogame.patchbot.patchbot as patchbot    
except ImportError:
    _LOG.warning(f"unable to import the patchbot. You won't be able to run the patching tests. Install from https://github.com/o-o-overflow/dcf-game-infrastructure")
try:
    import docker
except ImportError:
    _LOG.warning(f"unable to import docker. You won't be able to run the patching tests.")


SERVICE_DIR = os.path.dirname(__file__)

SERVICE_CONF = yaml.safe_load(open(os.path.join(SERVICE_DIR, "info.yml")))
SERVICE_NAME = SERVICE_CONF['service_name']
SERVICE_IMAGE = IMAGE_FMT % SERVICE_NAME
INTERACTION_IMAGE = IMAGE_FMT % SERVICE_NAME + '-interaction'
LOCAL_TESTER_IMAGE = IMAGE_FMT % SERVICE_NAME + '-local-tester'
DOCKER_REGISTRY = "registry.ctf:5000" if not 'DOCKER_REGISTRY' in os.environ else os.environ['DOCKER_REGISTRY']
NUM_TEAMS = 17

_LOG.info("USING YAML: %s/info.yml", SERVICE_DIR)
_LOG.info("SERVICE ID: %s", SERVICE_NAME)
_LOG.info("SERVICE IMAGE: %s", SERVICE_IMAGE)
_LOG.info("INTERACTION IMAGE: %s", INTERACTION_IMAGE)
_LOG.info("LOCAL TEST IMAGE: %s", LOCAL_TESTER_IMAGE)
_LOG.info("DOCKER REGISTRY: %s", DOCKER_REGISTRY)
_LOG.info("NUMBER OF TEAMS: %s", NUM_TEAMS)

@contextlib.contextmanager
def launch_container(image=SERVICE_IMAGE, container_id=None):
    _LOG.debug("launching container for image %s...", image)
    if container_id:
        assert os.system("docker run --name %s --rm -d -i %s" % (container_id, image)) == 0, "service container failed to start"
    else:
        container_id = subprocess.check_output("docker run --rm -d -i %s" % (image), shell=True).strip().decode('utf-8')
    _LOG.debug("container %s (image %s) launched!", container_id, image)
    time.sleep(1)

    # get endpoint info
    container_config = json.loads(subprocess.check_output("docker inspect %s" % container_id, shell=True).decode('utf-8'))
    ip_address = container_config[0]['NetworkSettings']['Networks']['bridge']['IPAddress']
    port = list(container_config[0]['Config']['ExposedPorts'].keys())[0].split("/")[0] if 'ExposedPorts' in container_config[0]['Config'] else None
    _LOG.debug("network endpoint: %s:%s", ip_address, port)

    test_patchable_variable(image, container_id)

    try:
        yield container_id, ip_address, port
    finally:
        _LOG.debug("stopping container %s", container_id)
        os.system("docker kill %s 2>/dev/null >/dev/null" % container_id)
        os.system("docker rm %s 2>/dev/null >/dev/null" % container_id)


def test_patchable_variable(image, container_id):
    # for service images of normal type verify the declared patchable file exist in the service container
    if image != SERVICE_IMAGE:
        _LOG.info("Skipping patchable tests for {img} b/c not a service image.".format(img=image))
        return
    if SERVICE_CONF['type'].lower() == "king_of_the_hill":
        _LOG.info("Skipping patchable tests for {img} b/c it is king of the hill.".format(img=image))
        return

    assert 'patchable_file' in SERVICE_CONF, "patchable_file key not found in yaml even though this is normal service."
    tmp_file = tempfile.mktemp()
    print("{} {}".format(os.path.isdir(tmp_file), os.path.isfile(tmp_file)))

    cmd = "docker cp {cid}:{pf} {tf}".format(cid=container_id, pf=SERVICE_CONF['patchable_file'], tf=tmp_file)
    _LOG.debug("running command {cmd}".format(cmd=cmd))
    assert os.system(cmd) == 0, "failed to copy patchable file from {cid} container using '{cmd}'".format(cid=container_id, cmd=cmd)

    assert os.path.exists(tmp_file), "{tf} does not exist from {pf}.".format(tf=tmp_file, pf=SERVICE_CONF['patchable_file'])
    try:
        assert os.path.isfile(tmp_file), "Copying {pf} resulted in a directory but only a single file is patchable".format(pf=SERVICE_CONF['patchable_file'])
        os.unlink(tmp_file)
    except AssertionError:
        shutil.rmtree(tmp_file)
        raise

    _LOG.info("Service passed patchable tests.")


def validate_yaml():
    _LOG.info("Validating yaml...")
    assert 'service_name' in SERVICE_CONF, "no service name specified"
    assert 'initial_flag' in SERVICE_CONF, "no service flag specified"
    if 'test flag' in SERVICE_CONF['initial_flag']: _LOG.critical("REMEBER TO CHANGE THE FLAG: %s looks like the test flag", SERVICE_CONF['initial_flag'])
    if not re.match(FLAG_RE, SERVICE_CONF['initial_flag']):
        _LOG.critical("FLAG %s DOES NOT CONFORM TO THE FLAG FORMAT", SERVICE_CONF['initial_flag'])
        if not SERVICE_CONF['violates_flag_format']:
            assert False, "you violated the flag format!"

    assert (not set(SERVICE_NAME) - set(string.ascii_lowercase + '-' + string.digits)) and SERVICE_NAME[0] in string.ascii_lowercase, "SERVICE NAME MUST BE KUBERNETES-COMPATIBLE (lowercase, '-', and digits, starting with a letter)"

    git_remote = subprocess.check_output(("git -C %s remote get-url origin"%SERVICE_DIR).split()).decode('latin1').split("/")[-1].split(".git")[0].strip()
    assert git_remote == CTF_TAG + "-" + SERVICE_NAME, "GIT REPOSITORY (%s) MUST BE THE SAME AS THE CTF TAG PLUS THE SERVICE NAME (%s-%s)" % (git_remote, CTF_TAG, SERVICE_NAME)

def build_images():
    if os.path.exists(os.path.join(SERVICE_DIR, "service", "Dockerfile")):
        _LOG.info("Building service image...")
        assert os.system("docker build -t %s %s/service" % (SERVICE_IMAGE, SERVICE_DIR)) == 0, "service docker image build failed"
    else:
        _LOG.warning("no dockerfile found for service...")

    if os.path.exists(os.path.join(SERVICE_DIR, "remote-interaction", "Dockerfile")):
        _LOG.info("Building interaction image...")
        assert os.system("docker build -t %s %s/remote-interaction" % (INTERACTION_IMAGE, SERVICE_DIR)) == 0, "interaction docker image build failed"
    else:
        _LOG.warning("no dockerfile found for remote interactions...")

    if os.path.exists(os.path.join(SERVICE_DIR, "local-tester", "Dockerfile")):
        _LOG.info("Building interaction image...")
        assert os.system("docker build -t %s --build-arg SERVICE=%s %s/local-tester" % (LOCAL_TESTER_IMAGE, SERVICE_IMAGE, SERVICE_DIR)) == 0, "local tester build failed"
    else:
        _LOG.warning("no dockerfile found for local tests...")

def test_local():
    if not os.path.exists(os.path.join(SERVICE_DIR, "local-tester", "Dockerfile")):
        return

    _LOG.info("Local-testing container...")
    with launch_container(image=LOCAL_TESTER_IMAGE) as (tester_container, _, _):
        _LOG.info("launching local tests")
        tests = SERVICE_CONF['local_tests']
        for script in tests:
            _LOG.info("launching %s", script)
            assert os.system("docker exec %s %s" % (tester_container, script)) == 0, "Check %s failed!" % script

def test_interactions():
    with launch_container() as (_, ip, port):
        run_interactions(ip, port, SERVICE_CONF['initial_flag'])

    if SERVICE_CONF['flag_path']:
        _LOG.info("testing interactions with random flag...")
        with launch_container() as (c, ip, port):
            newflag = "OOO{" + ''.join(random.choices(string.ascii_lowercase, k=16)) + "}"
            assert os.system("docker exec %s sh -c 'echo %s > %s'" % (c, newflag, SERVICE_CONF['flag_path'])) == 0
            run_interactions(ip, port, newflag)

    test_short_reads()

def run_interactions(ip_address, port, flag=SERVICE_CONF['initial_flag'], max_concurrency=SERVICE_CONF['concurrent_connections']):
    _LOG.info("Remote-testing container...")

    with launch_container(image=INTERACTION_IMAGE) as (interaction_container, _, _):
        def _run_interaction(script):
            _LOG.info("launching %s", script)

            try:
                output = subprocess.check_output("docker exec %s %s %s %s" % (interaction_container, script, ip_address, port), shell=True)
                if os.path.basename(script).startswith("exploit"):
                    _LOG.info("... checking for the flag")
                    assert flag.encode('utf-8') in output, "exploit %s did not print the flag"%script #pylint:disable=unsupported-membership-test
            except (AssertionError, subprocess.CalledProcessError):
                _LOG.error("check failed. Rerunning without stdio capture for convenience:")
                retval = os.system("docker exec %s %s %s %s" % (interaction_container, script, ip_address, port))
                _LOG.info("fyi, it returned %d that time", retval)
                raise


        _LOG.info("launching interaction scripts")
        interaction_files = SERVICE_CONF['remote_interactions']
        for f in interaction_files:
            _run_interaction(f)

        n = 2
        old_level = _LOG.level
        while n <= max_concurrency:
            _LOG.info("STRESS TESTING (%d concurrent connections)", n)
            _LOG.setLevel(max(logging.WARNING, old_level))
            with concurrent.futures.ThreadPoolExecutor(max_workers=n) as pool:
                results = pool.map(_run_interaction, (interaction_files*n)[:n])
            try:
                for _ in results:
                    pass
            except Exception as e: #pylint:disable=broad-except
                _LOG.error('One iteration returns an exception: %s', str(e), exc_info=True)
                sys.exit(1)

            _LOG.setLevel(old_level)

            n *= 2


def test_short_reads():
    with launch_container() as (service_container, ip_address, port):
        _LOG.info("SHORT-READ SANITY CHECK")
        ALLOWED_DIFF = 2
        start_num_procs = len(subprocess.check_output("docker exec %s ps aux" % service_container, shell=True).splitlines())

        assert os.system('docker run --rm ubuntu bash -ec "for i in {1..128}; do echo > /dev/tcp/%s/%s; done"' % (ip_address, port)) == 0
        _LOG.info("waiting for service to clean up after short reads")
        time.sleep(15)

        final_num_procs = len(subprocess.check_output("docker exec %s ps aux" % service_container, shell=True).splitlines())
        assert final_num_procs < (start_num_procs + ALLOWED_DIFF), "your service did not clean up after short reads starting procs = {sp} final={fp}".format(sp=start_num_procs, fp=final_num_procs)


def build_bundle(args=None):
    _LOG.info("building public bundle!")

    tempdir = tempfile.mkdtemp()
    public_path = os.path.join(tempdir, SERVICE_NAME)
    os.makedirs(public_path)
    for f in SERVICE_CONF['public_files']:
        _LOG.debug("copying file %s into public files", f)
        cmd = "cp -L %s/%s %s/%s" % (SERVICE_DIR, f, public_path, os.path.basename(f))
        print(os.getcwd(), cmd)
        assert os.system(cmd) == 0, "failed to retrieve public file %s" % f

    time.sleep(2)
    assert os.system("tar cvzf %s/public_bundle.tar.gz -C %s %s" % (SERVICE_DIR, tempdir, SERVICE_NAME)) == 0, "public file tarball failed; this should not be your fault"

    print("")
    print("")
    _LOG.critical("PLEASE VERIFY THAT THIS IS CORRECT: files in public bundle:")
    os.system("tar tvzf %s/public_bundle.tar.gz" % SERVICE_DIR)

    print("")
    print("")
    print("ATTENTION: PLEASE MAKE SURE THAT THE CONTENTS OF public_files.tar.gz (SHOWN ABOVE) MAKE SENSE.")
    print("")
    print("")

def tag(image, tag):
    tagcmd = "docker tag %s %s" % (image, tag)
    _LOG.info(tagcmd)
    assert os.system(tagcmd) == 0

def deploy_image(image):
    pushcmd = "docker push %s" % (image)
    _LOG.info(pushcmd)
    assert os.system(pushcmd) == 0

def tag_and_deploy(image, registry):
    tag(image, "%s/%s:latest" % (registry, image))
    deploy_image("%s/%s:latest" % (registry, image))

def build(args=None):
    build_images()
    build_bundle()

def test(args):
    if args.ip == None and args.port == None:
        test_local()
        test_interactions()
    else:
        run_interactions(args.ip, int(args.port), SERVICE_CONF['initial_flag'])

def launch(args):
    build_images()
    with launch_container(image=args.image) as (_, _ip_address, _port):
        print("")
        print("SERVICE RUNNING AT: %s %s" % (_ip_address, _port))
        print("nc %s %s" % (_ip_address, _port))
        print("./tester test %s %s" % (_ip_address, _port))
        print("%s:%s" % (_ip_address, _port))
        input()

def deploy(args=None):
    tag_and_deploy(SERVICE_IMAGE, DOCKER_REGISTRY)
    if SERVICE_CONF['type'].lower() == "normal":
        tag_and_deploy(INTERACTION_IMAGE, DOCKER_REGISTRY)
        if os.path.exists(os.path.join(SERVICE_DIR, "local-tester", "Dockerfile")):
            tag_and_deploy(LOCAL_TESTER_IMAGE, DOCKER_REGISTRY)
        for i in range(NUM_TEAMS):
            team_id = i+1
            _LOG.info("tagging for team %d" % (team_id))
            team_service_name = "%s/%s-team-%d:latest" % (DOCKER_REGISTRY, SERVICE_IMAGE, team_id)
            tag(SERVICE_IMAGE, team_service_name)
            deploy_image(team_service_name)

def get_expose_from_docker(repo_dir):
    ddata = open(os.path.join(repo_dir, "service", "Dockerfile"),"r").read()
    ptrn = re.compile(r"EXPOSE[ \s]{1,10}(?P<port>[0-9]{1,6})", flags=re.IGNORECASE)

    search = re.search(ptrn, ddata)

    if search:
        return int(search.group("port"))
    else:
        print("Error, port not in Dockerfile")
        exit(-1)


def patchbot_test_with_patch(patch_contents):
    client = docker.from_env()
    api = docker.APIClient()

    port = get_expose_from_docker(".")

    return patchbot.check_and_deploy_service(client,
                                             api,
                                             container_name=SERVICE_IMAGE,
                                             deployed_container_name=f"{SERVICE_IMAGE}-patch-test",
                                             remote_interaction_container_name=INTERACTION_IMAGE,
                                             local_interaction_container_name=LOCAL_TESTER_IMAGE,
                                             patch_id=1,
                                             service_id=41,
                                             service_port=port,
                                             remote_interaction_scripts=SERVICE_CONF['remote_interactions'],
                                             local_interaction_scripts=SERVICE_CONF['local_tests'],
                                             docker_location_to_patch=SERVICE_CONF['patchable_file'],
                                             new_file=patch_contents,
                                             max_bytes=SERVICE_CONF['max_patchable_bytes'],
                                             check_timeout=SERVICE_CONF.get('check_timeout', 120),
                                             pull_latest=False,
                                             deploy_service=False)
    

def test_all_patches(args=None):
    build_images()    
    logging.getLogger("patchbot").setLevel(logging.DEBUG)
    for patch_file in pathlib.Path(".").glob("patches/*"):
        with open(patch_file, 'rb') as patch:
            _, ext = os.path.splitext(patch_file)
            status = ext[1:]
            _LOG.info(f"going to test {patch_file}, expecting the result to be {status}")
            result, metadata = patchbot_test_with_patch(patch.read())
            _LOG.info(f"results from {patch_file}: {result} with public metadata: {metadata}")
            assert result.name.lower() == status.lower(), f"Testing patching failed for {patch_file}, expected {status} but got {result}"
            _LOG.info(f"All good, got expected results for {patch_file}")

def test_patch(args):
    build_images()    
    logging.getLogger("patchbot").setLevel(logging.DEBUG)
    patch_file = args.patch
    with open(patch_file, 'rb') as patch:
        result, metadata = patchbot_test_with_patch(patch.read())
        _LOG.info(f"results from {patch_file}: {result} with public metadata: {metadata}")

def do_all(args=None):
    build_images()
    test_local()
    test_interactions()
    if os.path.exists(os.path.join(SERVICE_DIR, "patches")):
        test_all_patches()
    build_bundle()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(prog="tester")
    subparsers = parser.add_subparsers(help="sub-command help")

    parser_bundle = subparsers.add_parser("bundle", help="bundle the service")
    parser_bundle.set_defaults(func=build_bundle)

    parser_build = subparsers.add_parser("build", help="build and bundle the service")
    parser_build.set_defaults(func=build)

    parser_test = subparsers.add_parser("test", help="test the service")
    parser_test.set_defaults(func=test)
    parser_test.add_argument('ip', help="ip address to test", default=None, nargs='?')
    parser_test.add_argument('port', help="port to test", default=None, nargs='?')

    parser_launch = subparsers.add_parser("launch", help="launch the service")
    parser_launch.set_defaults(func=launch)
    parser_launch.add_argument('image', help="image to launch", default=SERVICE_IMAGE, nargs='?')

    parser_deploy = subparsers.add_parser("deploy", help="deploy the service")
    parser_deploy.set_defaults(func=deploy)

    parser_test_patch = subparsers.add_parser("test-patch", help="Test a specific patch")
    parser_test_patch.add_argument('patch', help="file to try as patch")
    parser_test_patch.set_defaults(func=test_patch)

    parser_test_all_patches = subparsers.add_parser("test-all-patches", help="Test all patches in ./patches")
    parser_test_all_patches.set_defaults(func=test_all_patches)
    

    parser_deploy = subparsers.add_parser("all", help="Do everything")
    parser_deploy.set_defaults(func=do_all)

    _args = parser.parse_args()

    if _args == argparse.Namespace():
        do_all()
    else:
        _args.func(_args)