import paramiko, sys

ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect('140.44.4.21', username='root', password='dietpi', timeout=10)
print("[OK] Connected to 140.44.4.21 via SSH")

# 1. Fix DNS in /etc/resolv.conf and /boot/dietpi.txt
fix_dns_cmd = """
echo "nameserver 140.44.4.1" > /etc/resolv.conf
echo "nameserver 1.1.1.1" >> /etc/resolv.conf
sed -i 's/^AUTO_SETUP_NET_STATIC_DNS=.*/AUTO_SETUP_NET_STATIC_DNS=140.44.4.1 1.1.1.1/' /boot/dietpi.txt
getent hosts get.k3s.io
"""
stdin, stdout, stderr = ssh.exec_command(fix_dns_cmd)
print("DNS resolution test:", stdout.read().decode().strip())

# 2. Kill stuck whiptail / login process
kill_cmd = "pkill -9 whiptail || true; pkill -9 dietpi-login || true"
stdin, stdout, stderr = ssh.exec_command(kill_cmd)
print("Killed stuck first-run whiptail UI")

# 3. Inject SSH Key immediately
pub_key = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIOsRiwuL7nfm1Ruew+Xk+3tbzD9ONyKAO0q2QwkppEPZ snarfnet-automation-master@snarfnet.net"
ssh_inject = f"""
mkdir -p /root/.ssh /etc/dropbear
echo "{pub_key}" >> /root/.ssh/authorized_keys
echo "{pub_key}" >> /etc/dropbear/authorized_keys
chmod 700 /root/.ssh
chmod 600 /root/.ssh/authorized_keys /etc/dropbear/authorized_keys
"""
stdin, stdout, stderr = ssh.exec_command(ssh_inject)
print("[OK] Injected SSH automation master key into authorized_keys")

# 4. Run /boot/Automation_Custom_Script.sh
print("Executing /boot/Automation_Custom_Script.sh...")
stdin, stdout, stderr = ssh.exec_command("bash /boot/Automation_Custom_Script.sh")

# Stream output in real time
for line in iter(stdout.readline, ""):
    print(line, end="")

print("\nSTDERR:")
print(stderr.read().decode())

ssh.close()