#!/usr/bin/env python

import re
import sys
import json
import argparse
import urllib.parse

def generate_template(listen_address, listen_port, remote_host, remote_port):

    # we could all need that
    imports = [
        "os",
        "socket",
        "threading"
    ]

    partial_imports = {
        "hackingscripts.utils": ["util"],
        "hackingscripts.utils.packeter": ["Packer", "Parser"]
    }

    imports = "\n".join(f"import {i}" for i in sorted(imports, key=len))
    imports += "\n" + "\n".join(sorted(list(f"from {p} import {', '.join(i)}" for p, i in partial_imports.items()), key=len))
    return f"""#!/usr/bin/env python

#
# THE BASE OF THIS FILE WAS AUTOMATICALLY GENERATED BY {' '.join(sys.argv)}
# For more information, visit: https://git.romanh.de/Roman/HackingScripts
# 

{imports}

BUFFER_SIZE = 4096

class Packet:
    def __init__(self):
        pass

    @staticmethod
    def from_data(data):
        packet = Packet()
        parser = Parser(data)
        # TODO: auto-generated method stub
        return packet

    def pack(self):
        buf = Packer()
        # TODO: auto-generated method stub
        return buf.get()

def forward(source, destination):
    try:
        while True:
            data = source.recv(BUFFER_SIZE)
            if not data:
                break

            # TODO: Parse / Manipulate packet
            # packet = Packet.from_data(data)
            # repacked = packet.pack()

            destination.sendall(data)
    except Exception:
        pass
    finally:
        source.close()
        destination.close()

def handle_client(client_socket, remote_host, remote_port):
    try:
        remote_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        remote_socket.connect((remote_host, remote_port))
    except Exception as e:
        print(f"Failed to connect to remote: {{e}}")
        client_socket.close()
        return

    # Start bidirectional forwarding
    threading.Thread(target=forward, args=(client_socket, remote_socket), daemon=True).start()
    threading.Thread(target=forward, args=(remote_socket, client_socket), daemon=True).start()

def start_proxy(local_host, local_port, remote_host, remote_port):
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind((local_host, local_port))
    server.listen(100)

    print(f"[*] Forwarding from {{local_host}}:{{local_port}} to {{remote_host}}:{{remote_port}}")

    while True:
        client_socket, addr = server.accept()
        print(f"[+] Connection from {{addr[0]}}:{{addr[1]}}")
        threading.Thread(
            target=handle_client,
            args=(client_socket, remote_host, remote_port),
            daemon=True
        ).start()

if __name__ == "__main__":
    start_proxy({repr(listen_address)}, {listen_port}, {repr(remote_host)}, {remote_port})
"""

if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description="Exploit Template for tcp application attacks",
        formatter_class=argparse.RawTextHelpFormatter
    )

    parser.add_argument("la", type=str, help="Listen Address")
    parser.add_argument("lp", type=int, help="Listen Port", choices=range(1,65535+1))
    parser.add_argument("rh", type=str, help="Remote Host")
    parser.add_argument("rp", type=int, help="Remote Port", choices=range(1,65535+1))

    args = parser.parse_args()
    template = generate_template(args.la, args.lp, args.rh, args.rp)
    print(template)