import asyncio
import importlib
import traceback
import os
import argparse
import ipaddress
import socket

import forwarding
import handler


async def queue_handler(queue: asyncio.Queue):
    last_edit_timestamp = os.stat(handler.__file__).st_mtime_ns
    while True:
        tag, data = await queue.get()
        try:
            timestamp = os.stat(handler.__file__).st_mtime_ns
            if timestamp > last_edit_timestamp:
                importlib.reload(handler)
                last_edit_timestamp = timestamp

            if asyncio.iscoroutinefunction(handler.handle):
                await handler.handle(tag, data)
            else:
                handler.handle(tag, data)
        except Exception:
            traceback.print_exc()


async def main():
    p = argparse.ArgumentParser()
    p.add_argument('remote_host', type=str, help='Host you want to spy on')
    p.add_argument('remote_port', type=int, help='Port on remote_host')
    p.add_argument('--listen-host', type=str, default='0.0.0.0', help='MITM binds to this (default 0.0.0.0)')
    p.add_argument('--listen-port', type=int, default=1337, help='MITM port (default 1337)')
    p.add_argument('--output-host', type=str, default='10.227.12.40',
                   help='IP used to proxy remote_host (default 10.227.12.40)')
    p.add_argument('--output-port', type=int, default=1337, help='Port used by proxy (default 1337)')
    args = p.parse_args()

    remote_host_ip = ipaddress.IPv4Address(socket.gethostbyname(args.remote_host))
    listen_host_ip = ipaddress.IPv4Address(args.listen_host)
    output_host_ip = ipaddress.IPv4Address(args.output_host)

    # set up iptables

    os.system(f'sudo iptables '
              f'-t nat -A OUTPUT '
              f'-d {remote_host_ip} -p tcp --dport {args.remote_port} '
              f'-j DNAT --to 127.0.0.1:{args.listen_port}')

    os.system(f'sudo iptables '
              f'-t nat -A OUTPUT '
              f'-d {output_host_ip} -p tcp --dport {args.output_port} '
              f'-j DNAT --to {remote_host_ip}:{args.remote_port}')

    # run man-in-the-middle

    queue = asyncio.Queue()
    forwarder = forwarding.Forwarder(str(listen_host_ip), args.listen_port,
                                     str(output_host_ip), args.output_port,
                                     queue)

    queue_task = asyncio.create_task(queue_handler(queue))

    await forwarder.run()

    queue_task.cancel()

    # tear down iptables

    os.system(f'sudo iptables '
              f'-t nat -D OUTPUT '
              f'-d {remote_host_ip} -p tcp --dport {args.remote_port} '
              f'-j DNAT --to 127.0.0.1:{args.listen_port}')

    os.system(f'sudo iptables '
              f'-t nat -D OUTPUT '
              f'-d {output_host_ip} -p tcp --dport {args.output_port} '
              f'-j DNAT --to {remote_host_ip}:{args.remote_port}')


if __name__ == '__main__':
    asyncio.run(main())