import argparse
import requests
import socket
import yaml
from flask import Flask, Response, request
from typing import Any, Dict, Optional
import urllib.parse as urlparse

from bemani.protocol import EAmuseProtocol, Node

# Application configuration
app = Flask(__name__)
config: Dict[str, Any] = {}


def modify_request(config: Dict[str, Any], req_body: Node) -> Optional[Node]:
    # Not sure if there's any reason to modify requests, but its plumbed
    return None


def modify_response(config: Dict[str, Any], resp_body: Node) -> Optional[Node]:
    # Figure out if we need to modify anything in the response
    if resp_body.name != 'response':
        # Not what we expected, bail
        return None

    # Now we get to the meat of packet detection
    body = resp_body.children[0]

    # Is this a services packet? We need to modify this to point the rest of the
    # game packets at this proxy :(
    if body.name == 'services':
        for child in body.children:
            if child.name == 'item':
                if child.attribute('name') == 'ntp':
                    # Don't override this
                    continue
                elif child.attribute('name') == 'keepalive':
                    # Completely rewrite to point at the proxy server, otherwise if we're proxying
                    # to a local backend, we will end up giving out a local address and we will get
                    # Network NG.
                    address = socket.gethostbyname(config['keepalive'])
                    child.set_attribute(
                        'url',
                        f'http://{address}/core/keepalive?pa={address}&ia={address}&ga={address}&ma={address}&t1=2&t2=10',
                    )
                else:
                    # Get netloc to replace
                    url = urlparse.urlparse(child.attribute('url'))
                    defaultport = {
                        'http': 80,
                        'https': 443,
                    }.get(url.scheme, 0)
                    if config['local_port'] != defaultport:
                        new_url = child.attribute('url').replace(url.netloc, f'{config["local_host"]}:{config["local_port"]}')
                    else:
                        new_url = child.attribute('url').replace(url.netloc, f'{config["local_host"]}')
                    child.set_attribute('url', new_url)

        return resp_body

    return None


@app.route('/', defaults={'path': ''}, methods=['GET'])
@app.route('/<path:path>', methods=['GET'])
def receive_healthcheck(path: str) -> Response:
    if '*' in config['remote']:
        remote_host = config['remote']['*']['host']
        remote_port = config['remote']['*']['port']
    else:
        return Response("No route for default PCBID", 500)

    actual_path = f'/{path}'
    if request.query_string is not None and len(request.query_string) > 0:
        actual_path = actual_path + f'?{request.query_string.decode("ascii")}'

    # Make request to foreign service, using the same parameters
    r = requests.get(
        f'http://{remote_host}:{remote_port}{actual_path}',
        timeout=config['timeout'],
        allow_redirects=False,
    )

    headers = {}
    for header in ['Location']:
        if header in r.headers:
            headers[header] = r.headers[header]

    return Response(r.content, r.status_code, headers)


@app.route('/', defaults={'path': ''}, methods=['POST'])
@app.route('/<path:path>', methods=['POST'])
def receive_request(path: str) -> Response:
    # First, parse the packet itself
    client_proto = EAmuseProtocol()
    server_proto = EAmuseProtocol()
    remote_address = request.headers.get('X-Remote-Address', None)
    request_compression = request.headers.get('X-Compress', None)
    request_encryption = request.headers.get('X-Eamuse-Info', None)
    request_client = request.headers.get('User-Agent', None)

    actual_path = f'/{path}'
    if request.query_string is not None and len(request.query_string) > 0:
        actual_path = actual_path + f'?{request.query_string.decode("ascii")}'

    if config['verbose']:
        print(f"HTTP request for URI {actual_path}")
        print(f"Compression is {request_compression}")
        print(f"Encryption key is {request_encryption}")

    req = client_proto.decode(
        request_compression,
        request_encryption,
        request.data,
    )

    if req is None:
        # Nothing to do here
        return Response("Unrecognized packet!", 500)

    if config['verbose']:
        print("Original request to server:")
        print(req)

    # Grab PCBID for directing to mulitple servers
    pcbid = req.attribute('srcid')
    if pcbid in config['remote']:
        remote_host = config['remote'][pcbid]['host']
        remote_port = config['remote'][pcbid]['port']
    elif '*' in config['remote']:
        remote_host = config['remote']['*']['host']
        remote_port = config['remote']['*']['port']
    else:
        return Response(f"No route for PCBID {pcbid}", 500)

    modified_request = modify_request(config, req)
    if modified_request is None:
        # Return the original binary data instead of re-encoding it
        # to the exact same thing.
        req_binary = request.data
    else:
        if config['verbose']:
            print("Modified request to server:")
            print(modified_request)

        # Re-encode the modified packet
        req_binary = server_proto.encode(
            request_compression,
            request_encryption,
            modified_request,
            client_proto.last_text_encoding,
            client_proto.last_packet_encoding,
        )

    # Set up custom headers for remote request.
    headers = {
        # For lobby functionality, make sure the request receives
        # the original IP address
        'X-Remote-Address': remote_address or request.remote_addr,
        # Some remote servers can be somewhat buggy, so we make sure
        # to specify a range of encodings.
        'Accept-Encoding': 'identity, deflate, compress, gzip',
    }

    # Copy over required headers that are sent by game client.
    if request_compression:
        headers['X-Compress'] = request_compression
    else:
        headers['X-Compress'] = 'none'
    if request_encryption:
        headers['X-Eamuse-Info'] = request_encryption

    # Make sure to copy the user agent as well.
    if request_client is not None:
        headers['User-Agent'] = request_client

    # Make request to foreign service, using the same parameters
    prep_req = requests.Request(
        'POST',
        url=f'http://{remote_host}:{remote_port}{actual_path}',
        headers=headers,
        data=req_binary,
    ).prepare()
    sess = requests.Session()
    r = sess.send(prep_req, timeout=config['timeout'])

    if r.status_code != 200:
        # Failed on remote side
        return Response("Failed to get response!", 500)

    # Decode response, for modification if necessary
    response_compression = r.headers.get('X-Compress', None)
    response_encryption = r.headers.get('X-Eamuse-Info', None)
    resp = server_proto.decode(
        response_compression,
        response_encryption,
        r.content,
    )

    if resp is None:
        # Nothing to do here
        return Response("Unrecognized packet!", 500)

    if config['verbose']:
        print("Original response from server:")
        print(resp)

    modified_response = modify_response(config, resp)
    if modified_response is None:
        # Return the original response data instead of re-encoding it
        # to the exact same thing.
        resp_binary = r.content
    else:
        if config['verbose']:
            print("Modified response from server:")
            print(modified_response)

        # Re-encode the modified packet
        resp_binary = client_proto.encode(
            response_compression,
            response_encryption,
            modified_response,
        )

    # Some old clients are case sensitive, so be careful to capitalize
    # these responses here.
    flask_resp = Response(resp_binary)
    if response_compression is not None:
        flask_resp.headers['X-Compress'] = response_compression
    if response_encryption is not None:
        flask_resp.headers['X-Eamuse-Info'] = response_encryption
    return flask_resp


def load_proxy_config(filename: str) -> None:
    global config

    config_data = yaml.safe_load(open(filename))
    if 'pcbid' in config_data and config_data['pcbid'] is not None:
        for pcbid in config_data['pcbid']:
            remote_name = config_data['pcbid'][pcbid]
            remote_config = config_data['remote'][remote_name]
            config['remote'][pcbid] = remote_config


def load_config(filename: str) -> None:
    global config

    config_data = yaml.safe_load(open(filename))
    config.update({
        'local_host': config_data['local']['host'],
        'local_port': config_data['local']['port'],
        'verbose': config_data.get('verbose', False),
        'timeout': config_data.get('timeout', 30),
        'keepalive': config_data.get('keepalive', 'localhost'),
    })

    if 'default' in config_data:
        remote_config = config_data['remote'][config_data['default']]
        config.update({
            'remote': {
                '*': remote_config,
            },
        })
    else:
        config.update({
            'remote': {}
        })

    if 'pcbid' in config_data and config_data['pcbid'] is not None:
        load_proxy_config(filename)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="A utility to MITM non-SSL eAmusement connections.")
    parser.add_argument("-p", "--port", help="Port to listen on. Defaults to 9090", type=int, default=9090)
    parser.add_argument("-a", "--address", help="Address to listen on. Defaults to all addresses", type=str, default="0.0.0.0")
    parser.add_argument("-r", "--real-address", help="Real address we are listening on (for NAT and such)", type=str, default='127.0.0.1')
    parser.add_argument("-q", "--remote-port", help="Port to connect to.", type=int, required=True)
    parser.add_argument("-b", "--remote-address", help="Address to connect to.", type=str, required=True)
    parser.add_argument("-c", "--config", help="Configuration file for PCBID to remote server mapping.", type=str, default=None)
    parser.add_argument("-k", "--keepalive", help="Keepalive domain to advertise. Defaults to localhost", type=str, default='localhost')
    parser.add_argument("-v", "--verbose", help="Display verbose packet info.", action='store_true')
    parser.add_argument("-t", "--timeout", help="Timeout (in seconds) for proxy requests. Defaults to 30 seconds.", type=int, default=30)
    args = parser.parse_args()

    config.update({
        'local_host': args.real_address,
        'local_port': args.port,
        'remote': {
            '*': {
                'host': args.remote_address,
                'port': args.remote_port,
            },
        },
        'verbose': args.verbose,
        'timeout': args.timeout,
        'keepalive': args.keepalive,
    })

    # Fill in remote addresses for PCBIDs we should redirect to a non-default server
    if args.config is not None:
        load_proxy_config(args.config)

    app.run(host='0.0.0.0', port=args.port, debug=True)