#!/usr/bin/env python3
"""
Fully repair and harden sninfr2lnxwrk01 (.21) SD card mounted on .13 (/dev/sda):
1. Runs fsck.vfat on /dev/sda1
2. Runs e2fsck -fy on /dev/sda2
3. Sets `tune2fs -e continue /dev/sda2` so ext4 NEVER drops into read-only mode on bad sectors
4. Mounts /dev/sda2 on /mnt/sd_repair:
   - Verifies hostname (sninfr2lnxwrk01)
   - Sets strict permissions on /root (700) and /root/.ssh/authorized_keys (600)
   - Verifies unified SSH key in /root/.ssh/authorized_keys and /home/dietpi/.ssh/authorized_keys
   - Updates /etc/fstab with `errors=continue,noatime`
   - Configures /etc/systemd/timesyncd.conf pointing to 140.44.4.11 and 14.4.4.44
5. Cleanly syncs and unmounts the card.
"""
import os
import sys
import paramiko

sys.stdout.reconfigure(encoding="utf-8")

PRIV_KEY_PATH = os.path.expanduser("~/.ssh/id_ed25519_snarfnet")
PUB_KEY = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIOsRiwuL7nfm1Ruew+Xk+3tbzD9ONyKAO0q2QwkppEPZ snarfnet-automation-master@snarfnet.net"

TIMESYNCD_CONF = """[Time]
NTP=140.44.4.11 14.4.4.44
FallbackNTP=time.cloudflare.com time.google.com 0.us.pool.ntp.org 1.us.pool.ntp.org
PollIntervalMinSec=32
PollIntervalMaxSec=512
"""

REPAIR_SCRIPT = f"""#!/bin/bash
set -e

echo "=== 1. Verifying /dev/sda on .13 ==="
lsblk /dev/sda

echo "=== 2. Repairing VFAT boot partition /dev/sda1 ==="
umount /dev/sda1 2>/dev/null || true
fsck.vfat -a /dev/sda1 || true

echo "=== 3. Repairing ext4 root partition /dev/sda2 with e2fsck ==="
umount /dev/sda2 2>/dev/null || true
e2fsck -fy /dev/sda2 || true

echo "=== 4. Hardening ext4 filesystem with tune2fs -e continue ==="
tune2fs -e continue /dev/sda2
tune2fs -l /dev/sda2 | grep "Errors behavior"

echo "=== 5. Mounting /dev/sda2 to apply hardening and configuration ==="
mkdir -p /mnt/sd_repair
mount /dev/sda2 /mnt/sd_repair

if [ -f /mnt/sd_repair/etc/hostname ]; then
    echo "Confirmed target node hostname: $(cat /mnt/sd_repair/etc/hostname)"
fi

echo "=== 6. Fixing permissions and SSH keys ==="
chmod 700 /mnt/sd_repair/root
mkdir -p /mnt/sd_repair/root/.ssh
chmod 700 /mnt/sd_repair/root/.ssh
touch /mnt/sd_repair/root/.ssh/authorized_keys
chmod 600 /mnt/sd_repair/root/.ssh/authorized_keys

if ! grep -Fxq "{PUB_KEY}" /mnt/sd_repair/root/.ssh/authorized_keys; then
    echo "{PUB_KEY}" >> /mnt/sd_repair/root/.ssh/authorized_keys
    echo "  Injected key into /root/.ssh/authorized_keys"
else
    echo "  Key already present in /root/.ssh/authorized_keys"
fi

# Fix dietpi user if present
if [ -d /mnt/sd_repair/home/dietpi ]; then
    chmod 700 /mnt/sd_repair/home/dietpi
    mkdir -p /mnt/sd_repair/home/dietpi/.ssh
    chmod 700 /mnt/sd_repair/home/dietpi/.ssh
    touch /mnt/sd_repair/home/dietpi/.ssh/authorized_keys
    chmod 600 /mnt/sd_repair/home/dietpi/.ssh/authorized_keys
    if ! grep -Fxq "{PUB_KEY}" /mnt/sd_repair/home/dietpi/.ssh/authorized_keys; then
        echo "{PUB_KEY}" >> /mnt/sd_repair/home/dietpi/.ssh/authorized_keys
        echo "  Injected key into /home/dietpi/.ssh/authorized_keys"
    fi
fi

# Fix /etc/fstab to prevent read-only remounts
echo "=== 7. Updating /etc/fstab with errors=continue ==="
if [ -f /mnt/sd_repair/etc/fstab ]; then
    sed -i 's/errors=remount-ro/errors=continue/g' /mnt/sd_repair/etc/fstab
    sed -i '/[[:blank:]]\\/[[:blank:]]/s/defaults/defaults,errors=continue,noatime/' /mnt/sd_repair/etc/fstab 2>/dev/null || true
    echo "--- Current /etc/fstab ---"
    cat /mnt/sd_repair/etc/fstab
fi

# Configure timesyncd for internal NTP
echo "=== 8. Configuring Internal NTP Server in timesyncd.conf ==="
mkdir -p /mnt/sd_repair/etc/systemd
cat << 'EOF' > /mnt/sd_repair/etc/systemd/timesyncd.conf
{TIMESYNCD_CONF}
EOF

if [ -f /mnt/sd_repair/boot/dietpi.txt ]; then
    sed -i 's/^[[:blank:]]*CONFIG_NTP_MODE=.*/CONFIG_NTP_MODE=4/' /mnt/sd_repair/boot/dietpi.txt
fi

echo "=== 9. Syncing and cleanly unmounting ==="
sync
umount /mnt/sd_repair
rmdir /mnt/sd_repair
sync
echo "SUCCESS: SD card has been repaired, hardened against read-only remounts, and cleanly unmounted!"
"""

def main():
    pkey = paramiko.Ed25519Key.from_private_key_file(PRIV_KEY_PATH)
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    print("Connecting to 140.44.4.13...")
    client.connect("140.44.4.13", username="root", pkey=pkey, timeout=10)
    
    stdin, stdout, stderr = client.exec_command(f"bash -c '{REPAIR_SCRIPT}'")
    out = stdout.read().decode('utf-8', errors='replace')
    err = stderr.read().decode('utf-8', errors='replace')
    print(out)
    if err:
        print("ERR:\n", err)
    client.close()

if __name__ == "__main__":
    main()