deploy.py

#!/usr/bin/env python3

import argparse
import os
from utils import Colors, print_colored, print_banner, run_command, has_ssh_keys

# --- Configuration ---
PAYLOAD_DIR = os.path.join(os.path.dirname(__file__), "payloads")
REMOTE_TEMP_DIR = "~/pi_bootstrap_tmp"

def main():
    parser = argparse.ArgumentParser(description="Deploy bootstrap configuration to a Raspberry Pi.")
    parser.add_argument("target", help="The target in user@host format (e.g., pi@192.168.1.50)")
    parser.add_argument("--key", "-i", help="Path to the SSH private key file")
    args = parser.parse_args()

    target = args.target
    key_path = args.key
    
    print_banner("🚀 Pi Bootstrap", f"Target: {target}")

    # Check for SSH keys
    use_fancy_ui = has_ssh_keys(target, key_path)
    
    # Always use interactive mode for the setup step since it requires sudo password
    # But we can still use fancy UI for non-interactive steps
    if not use_fancy_ui:
        print_colored("⚠ SSH keys not detected. You will be asked for your password multiple times.", Colors.YELLOW)
        print_colored("Tip: Run 'ssh-copy-id user@host' in the future to skip this.\n", Colors.DIM)
    else:
        print_colored("ℹ️  The setup script will prompt for your sudo password on the Pi.\n", Colors.CYAN)

    # Define base commands
    ssh_cmd = ["ssh"]
    scp_cmd = ["scp"]
    if key_path:
        ssh_cmd.extend(["-i", key_path])
        scp_cmd.extend(["-i", key_path])

    # Define steps - split interactive from non-interactive
    non_interactive_steps = [
        ("Removing potential conflicting file or directory...", ssh_cmd + [target, "rm -rf ~/repos/ohmyzsh"]),
        ("Creating temporary directory...", ssh_cmd + [target, f"mkdir -p {REMOTE_TEMP_DIR}"]),
        ("Uploading configuration files...", scp_cmd + ["-r", PAYLOAD_DIR, f"{target}:{REMOTE_TEMP_DIR}"])
    ]
    
    interactive_step = ("Running setup script...", ssh_cmd + ["-t", target, f"cd {REMOTE_TEMP_DIR}/payloads && chmod +x setup.sh && ./setup.sh"])
    
    cleanup_step = ("Cleaning up...", ssh_cmd + [target, f"rm -rf {REMOTE_TEMP_DIR}"])

    if use_fancy_ui:
        # Run non-interactive steps
        for desc, cmd in non_interactive_steps:
            print_colored(f"➜ {desc}", Colors.CYAN)
            run_command(cmd, shell=isinstance(cmd, str))
            print_colored(f"  ✔ {desc.replace('...', '')}", Colors.GREEN)
        
        # Run interactive step
        desc, cmd = interactive_step
        print_colored(f"➜ {desc}", Colors.CYAN)
        run_command(cmd, shell=isinstance(cmd, str), interactive=True)
        print_colored(f"  ✔ {desc.replace('...', '')}", Colors.GREEN)
        
        # Run cleanup
        desc, cmd = cleanup_step
        print_colored(f"➜ {desc}", Colors.CYAN)
        run_command(cmd, shell=isinstance(cmd, str))
        print_colored(f"  ✔ {desc.replace('...', '')}", Colors.GREEN)
    else:
        # Interactive mode (show all output)
        all_steps = non_interactive_steps + [interactive_step, cleanup_step]
        for desc, cmd in all_steps:
            print_colored(f"➜ {desc}", Colors.CYAN)
            run_command(cmd, shell=isinstance(cmd, str), interactive=True)
            print_colored(f"  ✔ Done\n", Colors.GREEN)

    print_colored(f"\n✨ Bootstrap Successful!", f"{Colors.BOLD}{Colors.GREEN}")
    print(f"Please log out and log back in to {Colors.BOLD}{target}{Colors.RESET} to see the changes.\n")

if __name__ == "__main__":
    main()