#!/usr/bin/env python3
"""
Fully repair and harden sninfr2lnxwrk01 (.21) SD card mounted on .22:
1. Detects the inserted SD card reader (e.g. /dev/sdd or /dev/sdX).
2. Runs e2fsck -fy on the ext4 root partition to repair all corruption and orphan inodes.
3. Sets `tune2fs -e continue` so ext4 NEVER drops into read-only mode on bad sectors.
4. Mounts the partition:
   - Sets strict permissions on /root (700) and /root/.ssh/authorized_keys (600).
   - Injects the unified SSH key if missing.
   - Configures /etc/fstab with `errors=continue,noatime`.
   - Configures /etc/systemd/timesyncd.conf pointing to 140.44.4.11 and 14.4.4.44.
   - Enables continuous NTP daemon mode if DietPi.
5. Cleanly syncs and unmounts the card.
"""
import os
import sys
import time
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. Detecting SD card reader on .22 ==="
# Find disk that is removable or matches SDDR / USB reader, size ~238G or not mmcblk0/sda/sdb/sdc
SD_DEV=""
for dev in /sys/block/sd*; do
    DEVNAME=$(basename "$dev")
    # skip sda (Flash Drive FIT 128G) and sdb/sdc (Longhorn 8G virtual disks)
    if [ "$DEVNAME" = "sda" ] || [ "$DEVNAME" = "sdb" ] || [ "$DEVNAME" = "sdc" ]; then
        continue
    fi
    if [ -b "/dev/$DEVNAME" ]; then
        SD_DEV="/dev/$DEVNAME"
        echo "Found SD reader device: $SD_DEV"
        break
    fi
done

if [ -z "$SD_DEV" ]; then
    echo "ERROR: No candidate SD reader found on .22! Please plug in the USB adapter."
    lsblk -o NAME,SIZE,FSTYPE,MODEL
    exit 1
fi

ROOT_PART="${{SD_DEV}}2"
BOOT_PART="${{SD_DEV}}1"

if [ ! -b "$ROOT_PART" ]; then
    echo "ERROR: Partition $ROOT_PART does not exist!"
    lsblk "$SD_DEV"
    exit 1
fi

echo "=== 2. Running e2fsck on $ROOT_PART ==="
# Unmount if mounted anywhere
umount "$ROOT_PART" 2>/dev/null || true
e2fsck -fy "$ROOT_PART" || true

echo "=== 3. Hardening ext4 filesystem with tune2fs -e continue ==="
tune2fs -e continue "$ROOT_PART"
tune2fs -l "$ROOT_PART" | grep "Errors behavior"

echo "=== 4. Mounting $ROOT_PART to apply configuration ==="
mkdir -p /mnt/sd_repair
mount "$ROOT_PART" /mnt/sd_repair

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

# Fix /root and /root/.ssh permissions
echo "=== 5. 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 "=== 6. 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 "=== 7. 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 "=== 8. 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.22...")
    client.connect("140.44.4.22", 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()