import paramiko

def apply_to_host(ip):
    print(f"=== Applying to {ip} ===")
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(ip, username='root', key_filename=r'C:\Users\Snarf\.ssh\id_ed25519_snarfnet', timeout=10)

    commands = [
        "iptables -t nat -C PREROUTING -p tcp --dport 10051 -j DNAT --to-destination 140.44.4.64:10051 2>/dev/null || iptables -t nat -I PREROUTING 1 -p tcp --dport 10051 -j DNAT --to-destination 140.44.4.64:10051",
        "iptables -t nat -C POSTROUTING -p tcp -d 140.44.4.64 --dport 10051 -j MASQUERADE 2>/dev/null || iptables -t nat -I POSTROUTING 1 -p tcp -d 140.44.4.64 --dport 10051 -j MASQUERADE",
        "iptables -t nat -C OUTPUT -p tcp -d 140.44.4.70 --dport 10051 -j DNAT --to-destination 140.44.4.64:10051 2>/dev/null || iptables -t nat -I OUTPUT 1 -p tcp -d 140.44.4.70 --dport 10051 -j DNAT --to-destination 140.44.4.64:10051",
        "iptables -t nat -C OUTPUT -p tcp -d 140.44.4.71 --dport 10051 -j DNAT --to-destination 140.44.4.64:10051 2>/dev/null || iptables -t nat -I OUTPUT 2 -p tcp -d 140.44.4.71 --dport 10051 -j DNAT --to-destination 140.44.4.64:10051",
        "iptables -t nat -C OUTPUT -p tcp -d 140.44.4.72 --dport 10051 -j DNAT --to-destination 140.44.4.64:10051 2>/dev/null || iptables -t nat -I OUTPUT 3 -p tcp -d 140.44.4.72 --dport 10051 -j DNAT --to-destination 140.44.4.64:10051",
        "iptables -C FORWARD -m state --state RELATED,ESTABLISHED -j ACCEPT 2>/dev/null || iptables -I FORWARD 1 -m state --state RELATED,ESTABLISHED -j ACCEPT",
        "iptables -C FORWARD -p tcp -d 140.44.4.64 --dport 10051 -j ACCEPT 2>/dev/null || iptables -I FORWARD 2 -p tcp -d 140.44.4.64 --dport 10051 -j ACCEPT",
        "sysctl -w net.ipv4.ip_forward=1",
        "which iptables-save"
    ]

    for cmd in commands:
        _, o, e = ssh.exec_command(cmd)
        out = o.read().decode().strip()
        if out:
            print(f"  {out}")
        err = e.read().decode().strip()
        if err:
            print(f"  [err] {err}")

    # Check persistence package (iptables-persistent or netfilter-persistent)
    stdin, stdout, stderr = ssh.exec_command("which netfilter-persistent || which iptables-save")
    print(f"  Persistence tool: {stdout.read().decode().strip()}")
    stdin, stdout, stderr = ssh.exec_command("iptables-save > /etc/iptables/rules.v4 2>/dev/null || iptables-save > /etc/iptables.rules 2>/dev/null || true")

    ssh.close()

apply_to_host("140.44.4.71")
apply_to_host("140.44.4.72")