# i2p-docker-proxy
# Copyright (C) 2019 LoveIsGrief
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
import argparse
import asyncio
import logging
import multiprocessing
import os
import typing
from time import sleep


from i2plib import sam

from trans_proxy import fake_dns
from trans_proxy.process import AsyncProcess
from trans_proxy.servers import ClientTcpTunnel

ENV_PORT = "PROXY_PORT"
ENV_SAM_HOST = "PROXY_SAM_HOST"
ENV_SAM_PORT = "PROXY_SAM_PORT"
ENV_DNS_PORT = "PROXY_DNS_PORT"

logger = logging.getLogger("trans_proxy")


def main():
    parser = argparse.ArgumentParser(
        description="Forwards packets to an I2P instance and handles DNS requests"
    )
    parser.add_argument(
        "--verbose",
        action='store_true',
        help="Activates verbose logs")
    parser.add_argument(
        "-p", "--port",
        default=os.environ.get(ENV_PORT, 1234),
        type=int,
        help="Where all traffic should enter to be forwarded")
    parser.add_argument(
        "--sam-host",
        default=os.environ.get(ENV_SAM_HOST, sam.DEFAULT_ADDRESS[0]))
    parser.add_argument(
        "--sam-port",
        type=int,
        default=os.environ.get(ENV_SAM_PORT, sam.DEFAULT_ADDRESS[1]))

    dns_group = parser.add_argument_group('dns')
    dns_group.add_argument(
        "--dns-port",
        type=int,
        default=os.environ.get(ENV_DNS_PORT, 1053))

    args = parser.parse_args()

    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)

    with multiprocessing.Manager() as m:
        ip_dict = m.dict()
        processes = [
            AsyncProcess(
                target=start_client_tcp_tunnel,
                kwargs={
                    "sam_host": args.sam_host,
                    "sam_port": args.sam_port,
                    "ip_dict": ip_dict,
                    "port": args.port,
                }
            ),
            multiprocessing.Process(
                target=fake_dns.main,
                kwargs={
                    "port": args.dns_port,
                    "ip_dict": ip_dict,
                }
            )
        ]
        exec_processes(processes)


def exec_processes(processes: typing.List[multiprocessing.Process]):
    """
    Starts processes and waits for them to end
    """
    # TODO: manage processes better
    for process in processes:
        process.start()

    try:
        all_processes_done = False
        while not all_processes_done:
            sleep(1)
            all_processes_done = all(
                not process.is_alive()
                for process in processes
            )
    except KeyboardInterrupt:
        pass
    finally:
        for process in processes:
            if process.is_alive():
                process.close()


async def start_client_tcp_tunnel(
        sam_host,
        sam_port,
        ip_dict,
        host="127.0.0.1", port=1234,
        **kwargs
):
    loop = asyncio.get_running_loop()
    server = await loop.create_server(lambda: ClientTcpTunnel(
        sam_host=sam_host,
        sam_port=sam_port,
        ip_dict=ip_dict,
    ), host=host, port=port)
    await server.serve_forever()


if __name__ == '__main__':
    main()