from pwn import log, remote, process
import sys
from argparse import ArgumentParser


def connect_tcp(local_path):
    # Argument parsing
    parser = ArgumentParser()
    parser.add_argument("--connection-info", required=True, help="nc HOST PORT")
    parser.add_argument("--local", action="store_true", help="Run locally")
    args = parser.parse_args()

    # Parse HOST and PORT
    try:
        _, HOST, PORT = args.connection_info.split(" ")
        PORT = int(PORT)
    except ValueError:
        log.error("Invalid --connection-info format. Use 'nc HOST PORT'")
        sys.exit(1)

    # Connect to process or remote host
    conn = process(["python", local_path]) if args.local else remote(HOST, PORT)
    return conn