server.py

# This file contains the server which maintains a websocket persistant connexion
# with the client.

# It also is a HTTP server to recieve the HTTP requests from the proxy
# and convert them to orders to the client

import asyncio
import websockets
from flask import Flask, request, Response, render_template
from werkzeug.routing import Rule
import threading
import random
import string
import json
from proxy import start_proxy, MITMparams

# Config part

config = {}
mitm_params = MITMparams()

# Websocket part

CLIENTS = {}

def generate_random_string(length):
    letters = string.ascii_lowercase
    return ''.join(random.choice(letters) for i in range(length))

def get_client_id():
    # Return a ID of 10 character, which is not currently set inside CLIENTS
    while True:
        result_str = generate_random_string(10)
        if result_str not in CLIENTS:
            break
    return result_str

async def handle_connexion(websocket):
    id_connection = get_client_id()
    CLIENTS[id_connection] = {"ws": websocket, "commands": {}, "origin": websocket.origin}
    set_clients()
    try:
        async for _ in websocket:
            if (_):
                resp = json.loads(_)
                CLIENTS[id_connection]["commands"][resp["id_request"]] = resp

            pass
    except Exception as e:
        print(f"Exception in websocket connexion : {e}", flush=True)
    finally:
        del CLIENTS[id_connection]
        set_clients()
        print(f"Client {id_connection} disconnected", flush=True)
    
async def websockets_server():
    async with websockets.serve(handle_connexion, config["websocket_server"]["interface"], config["websocket_server"]["port"], max_size=3000000000):
        await asyncio.Future()  # run forever

async def sendCommand(websocket_client_id, command):
    try:
        CLIENTS[websocket_client_id]["response"] = ""
        await CLIENTS[websocket_client_id]["ws"].send(json.dumps(command))
        while command["id_request"] not in CLIENTS[websocket_client_id]["commands"]:
            pass
        resp = CLIENTS[websocket_client_id]["commands"][command["id_request"]]

        del CLIENTS[websocket_client_id]["commands"][command["id_request"]]
        return resp

    except websockets.ConnectionClosed:
        pass

def run_websocket_server():
    asyncio.run(websockets_server())

# Webserver part

app = Flask(__name__, template_folder="html/templates")

# The following line allow the usage of any HTTP methods
app.url_rule_class = lambda rule, **kwargs: Rule(rule, **{**kwargs, 'methods': None})

@app.route('/', defaults={'path': ''})
@app.route("/")
def default_route():
    return render_template("main.html", sessions=CLIENTS)


@app.route('/<path:path>')
async def proxy_route(path):
    proxified_request = {
        "id_request": generate_random_string(20),
        "url": request.cookies.get('full_url'),
        "method": request.method,
        "sendRequest": True,
        "data": request.get_data().decode('utf-8'),
        "headers": dict(request.headers)
    }
    websocket_client = path
    try:
        proxified_response = await sendCommand(websocket_client_id=websocket_client, command=proxified_request)
        if "headers" in proxified_response:
            headers = proxified_response["headers"]
        else:
            headers = {}
        if "data" not in proxified_response:
            resp = Response({})
        else:
            resp = Response(proxified_response["data"])
        tunnelled_headers = {}
        tunnelled_headers["mitmproxy_override"] = json.dumps(headers)
        resp.headers = tunnelled_headers
        return resp
    except Exception as e:
        print(f"Exception in route proxy_route : {e}", flush=True)

# Proxy part

def start_proxy_server():
    asyncio.run(start_proxy(mitm_params, config))

def get_whitelisted_domains():
    if config["web_server"]["ip_address"] in ["127.0.0.1", "localhost"]:
        whitelisted_urls = [
            f'127.0.0.1:{config["web_server"]["port"]}',
            f'localhost:{config["web_server"]["port"]}'
        ]
    else:
        whitelisted_urls = [f'{config["web_server"]["ip_address"]}:{config["web_server"]["port"]}']
    if config["web_server"]["port"] == 443:
        for i in whitelisted_urls:
            whitelisted_urls.append(i.split(':')[0])
    return whitelisted_urls

def set_clients():
    result = {}
    for i in CLIENTS:
        result[CLIENTS[i]["origin"]] = i
    mitm_params.set_clients(result)

# Main part

if __name__ == "__main__":
    try:
        with open('config.json') as config_file:
            file_contents = config_file.read()

        config = json.loads(file_contents)
        
        # websocket
        websocket_thread = threading.Thread(target=run_websocket_server)
        websocket_thread.daemon = True
        websocket_thread.start()

        # proxy
        config["proxy"]["do_not_proxify"]["domains"] = get_whitelisted_domains()
        mitm_params.set_config(config)
        proxy_thread = threading.Thread(target=start_proxy_server)
        proxy_thread.daemon = True
        proxy_thread.start()

        # webserver
        app.run(ssl_context=('./certificates/CA_webserver.crt', './certificates/cakey_webserver.pem'), port=config["web_server"]["port"], host=config["web_server"]["interface"])
    except (KeyboardInterrupt, SystemExit):
        print('Received keyboard interrupt, quitting threads.', flush=True)