#!/usr/bin/env python3

from sanic import Sanic
from sanic.response import html, redirect, text, raw
from sanic.exceptions import NotFound
import sys
import asyncio
import ssl
from pathlib import Path
import requests
import json
import os
import aiohttp
from OpenSSL import crypto, SSL
import random

config = {}

async def print_flush(message):
    print(message)
    sys.stdout.flush()

session = aiohttp.ClientSession()

def server_error_handler(request, exception):
    if request is None:
        print("Invalid HTTP Method - Likely HTTPS trying to talk to HTTP")
        sys.stdout.flush()
        return html("Error: Failed to process request", status=500, headers={})
    return html("Error: Requested URL {} not found".format(request.url), status=404, headers=config[request.app.name]['headers'])


async def download_file(request, **kwargs):
    try:
        if config[request.app.name]['debug']:
            await print_flush("agent_message request from: {} with {} and {}".format(request.url, request.cookies, request.headers))
            await print_flush(config[request.app.name]['payloads'])
            await print_flush(request.path)
        if config[request.app.name]['debug']:
            await print_flush(f"forwarding to : {config['mythic_base_address'] + '/api/v1.4/files/download/{}'.format(config[request.app.name]['payloads'][request.path])}")
        async with session.get(config['mythic_base_address'] + "/api/v1.4/files/download/{}".format(config[request.app.name]['payloads'][request.path]), ssl=False, headers={"Mythic": "http", **request.headers}) as resp:
            return raw(await resp.read(), status=resp.status, headers=config[request.app.name]['headers'])
    except Exception as e:
        await print_flush(str(e))

async def agent_message(request, **kwargs):
    global config
    forwarded_headers = {
        "x-forwarded-for": request.headers["x-forwarded-for"] if "x-forwarded-for" in request.headers else request.ip,
        "x-forwarded-url": request.url,
        "x-forwarded-query": request.query_string,
        "x-forwarded-port": str(request.server_port),
        "x-forwarded-cookies": str(request.cookies)
    }
    try:
        if config[request.app.name]['debug']:
            await print_flush("agent_message request from: {} with {} and {}".format(request.url, request.cookies, request.headers))
            await print_flush(" and URI: {}".format(request.query_string))
        if config[request.app.name]['debug']:
            await print_flush("Forwarding along to: {}".format(config['mythic_address']))
        if request.method == "POST":
            # manipulate the request if needed
            #from mythic_c2_container.MythicRPC import MythicRPC
            #msg = await MythicRPC().execute("create_encrypted_message", message={"action": "get_tasking"}, target_uuid="ba3f9b0a-e073-4cbc-afe7-1c94f86cc76a", c2_profile_name="http")
            #resp = await MythicRPC().execute("create_event_message", message=msg.response, warning=False)
            #msg = await MythicRPC().execute("create_decrypted_message", message=msg.response, c2_profile_name="http")
            #resp = await MythicRPC().execute("create_event_message", message=msg.response, warning=False)
            async with session.post(config['mythic_address'], data=request.body,  ssl=False, headers={"Mythic": "http", **request.headers, **forwarded_headers}) as resp:
                return raw(await resp.read(), status=resp.status, headers=config[request.app.name]['headers'])
        else:
            # manipulate the request if needed
            async with session.get(config['mythic_address'] + "?{}".format(request.query_string), data=request.body, ssl=False, headers={"Mythic": "http", **request.headers, **forwarded_headers}) as resp:
                return raw(await resp.read(), status=resp.status, headers=config[request.app.name]['headers'])
    except Exception as e:
        if request is None:
            await print_flush("Invalid HTTP Method - Likely HTTPS trying to talk to HTTP")
            return server_error_handler(request, e)
        if config[request.app.name]['debug']:
            await print_flush("error in agent_message: {}".format(str(e)))
        return server_error_handler(request, e)

def random_mixed():
    random.seed()
    length = random.randint(5,20)
    letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
    rnd = ''.join(random.choice(letters) for i in range(length))
    return rnd

# from https://stackoverflow.com/questions/27164354/create-a-self-signed-x509-certificate-in-python
def cert_gen(
    emailAddress: str = random_mixed(),
    commonName: str = random_mixed(),
    countryName: str = "US",
    localityName: str = random_mixed(),
    stateOrProvinceName: str = random_mixed(),
    organizationName: str = random_mixed(),
    organizationUnitName: str = random_mixed(),
    serialNumber: int = 0,
    validityStartInSeconds: int = 0,
    validityEndInSeconds: int = 10*365*24*60*60,
    key_file: str  = "privkey.pem",
    cert_file: str = "fullchain.pem"):

    #can look at generated file using openssl:
    #openssl x509 -inform pem -in selfsigned.crt -noout -text
    # create a key pair
    k = crypto.PKey()
    k.generate_key(crypto.TYPE_RSA, 4096)
    # create a self-signed cert
    cert = crypto.X509()
    cert.get_subject().C = countryName
    cert.get_subject().ST = stateOrProvinceName
    cert.get_subject().L = localityName
    cert.get_subject().O = organizationName
    cert.get_subject().OU = organizationUnitName
    cert.get_subject().CN = commonName
    cert.get_subject().emailAddress = emailAddress
    cert.set_serial_number(serialNumber)
    cert.gmtime_adj_notBefore(0)
    cert.gmtime_adj_notAfter(validityEndInSeconds)
    cert.set_issuer(cert.get_subject())
    cert.set_pubkey(k)
    cert.sign(k, 'sha512')
    with open(cert_file, "wt") as f:
        f.write(crypto.dump_certificate(crypto.FILETYPE_PEM, cert).decode("utf-8"))
    with open(key_file, "wt") as f:
        f.write(crypto.dump_privatekey(crypto.FILETYPE_PEM, k).decode("utf-8"))



if __name__ == "__main__":
    sys.path.append("/Mythic/mythic")
    config_file = open("config.json", 'rb')
    main_config = json.loads(config_file.read().decode('utf-8'))
    print("Opening config and starting instances...")
    sys.stdout.flush()
    # basic mapping of the general endpoints to the real endpoints
    try:
        config['mythic_address'] = os.environ['MYTHIC_ADDRESS']
        config['mythic_base_address'] = config['mythic_address'].split('/api')[0]
    except Exception as e:
        print("failed to find MYTHIC_ADDRESS environment variable")
        sys.stdout.flush()
        sys.exit(1)
    # now look at the specific instances to start
    for inst in main_config['instances']:
        config[str(inst['port'])] = {'debug': inst['debug'],
                                     'headers': inst['ServerHeaders'],
                                     'payloads': {}}
        if inst['debug']:
            print("Debugging statements are enabled. This gives more context, but might be a performance hit")
        else:
            print("Debugging statements are disabled")
        sys.stdout.flush()
        # now to create an app instance to handle responses
        app = Sanic(str(inst['port']))
        app.config['REQUEST_MAX_SIZE'] = 1000000000
        app.config['REQUEST_TIMEOUT'] = 600
        app.config['RESPONSE_TIMEOUT'] = 600
        if "payloads" in inst and isinstance(inst["payloads"], dict):
            for k, v in inst["payloads"].items():
                config[str(inst["port"])]["payloads"][f"{k}"] = v
                app.add_route(download_file, f"{k}", methods=["GET"])
        app.add_route(agent_message, "/<uri:path>", methods=['GET','POST'])
        app.add_route(agent_message, "/", methods=['GET','POST'])
        app.error_handler.add(Exception, server_error_handler)
        keyfile = Path(inst['key_path'])
        certfile = Path(inst['cert_path'])
        if (not keyfile.is_file() or not certfile.is_file()) and inst["use_ssl"]:
            # we need to generate the certs
            cert_gen(key_file=inst['key_path'], cert_file=inst['cert_path'])
        if keyfile.is_file() and certfile.is_file() and inst['use_ssl']:
            context = ssl.create_default_context(purpose=ssl.Purpose.CLIENT_AUTH)
            context.load_cert_chain(inst['cert_path'], keyfile=inst['key_path'])
            if inst['debug']:
                server = app.create_server(host="0.0.0.0", port=inst['port'], ssl=context, debug=False, return_asyncio_server=True, access_log=True)
            else:
                server = app.create_server(host="0.0.0.0", port=inst['port'], ssl=context, debug=False, return_asyncio_server=True, access_log=False)
            print("++++++++++++++++++++++++++++++++++++++++++")
            print("++++++++++++++++++++++++++++++++++++++++++")
            print("Successfully using SSL for port {}".format(inst['port']))
            print("++++++++++++++++++++++++++++++++++++++++++")
            print("++++++++++++++++++++++++++++++++++++++++++")
            sys.stdout.flush()
        else:
            if (inst["key_path"] != "" or inst["cert_path"] != "") and inst["use_ssl"]:
                print("-----------------------------------------------------------------------------")
                print("-----------------------------------------------------------------------------")
                print("key_path or cert_path is not empty and use_ssl is true, but at least one of the files cannot be found.")
                print("If you're using the http profile in a Docker container, place your cert and key in Mythic/C2_Profiles/http/c2_code/")
                print("    and set the key_path and cert_path values to just the names of the files")
                print("If you're using the http profile outside of Docker, make sure the paths are absolute and correct")
                print("NOT using SSL for port {}".format(inst['port']))
                print("-----------------------------------------------------------------------------")
                print("-----------------------------------------------------------------------------")
                sys.stdout.flush()
            else:
                print("**************************************************************")
                print("**************************************************************")
                print("NOT using SSL for port {}".format(inst['port']))
                print("**************************************************************")
                print("**************************************************************")
                sys.stdout.flush()
            if inst['debug']:
                server = app.create_server(host="0.0.0.0", port=inst['port'], debug=False, return_asyncio_server=True, access_log=True)
            else:
                server = app.create_server(host="0.0.0.0", port=inst['port'], debug=False, return_asyncio_server=True, access_log=False)
        task = asyncio.ensure_future(server)

    try:
        loop = asyncio.get_event_loop()
        def callback(fut):
            try:
                fetch_count = fut.result()
            except:
                print("port already in use")
                sys.stdout.flush()
                sys.exit()
        task.add_done_callback(callback)
        loop.run_forever()
    except:
        sys.exit()
        loop.stop()
