Draft version of Strato Pi upgrade script (python script orchestrating bash subscripts)

stratopi_upgrade.py

#!/usr/bin/env python3

import subprocess
import tempfile
import os
import json
import argparse
from datetime import datetime
from pathlib import Path
import shutil
import glob

# Default configuration
DEFAULT_CONFIG = {
    "services": [
        "mongodb.service"
    ],
    "kernel_modules": [
        "stratopi",
        "rtc_ds1307"
    ],
    "cleanup_paths": [
        "~/module_backup",
        "~/deployments/*"
    ],
    "remote_server": "192.0.2.10",
    "remote_port": 2222,
    "remote_user": "logs_user",
    "remote_path": "/var/log/rpi-upgrades",
    "ssh_pass": "secure_password"
}

LOG_DIR = Path("/home/pi/.rpi_upgrade/logs")
LOG_DIR.mkdir(parents=True, exist_ok=True)


def load_config(config_path):
    if config_path and Path(config_path).is_file():
        with open(config_path, 'r') as f:
            user_config = json.load(f)
        return {**DEFAULT_CONFIG, **user_config}
    return DEFAULT_CONFIG


def run_script(script_content, log_filename):
    with tempfile.NamedTemporaryFile('w+', delete=False, suffix='.sh') as script:
        script.write(script_content)
        script.flush()
        script_path = script.name

    os.chmod(script_path, 0o755)

    with open(LOG_DIR / log_filename, 'w') as logfile:
        subprocess.run(["sudo", "bash", script_path], stdout=logfile, stderr=subprocess.STDOUT)

    os.unlink(script_path)


def pre_upgrade_script(services, kernel_modules, cleanup_paths):
    unload_modules = "\n".join([f"modprobe -r {mod}" for mod in kernel_modules])
    stop_services = "\n".join([
        f"systemctl stop {svc} || service {svc} stop; systemctl disable {svc}; systemctl mask {svc}; rm -f /etc/systemd/system/{svc}"
        for svc in services
    ])
    cleanup_cmds = "\n".join(
        [f"rm -rf {path}" for pattern in cleanup_paths for path in glob.glob(os.path.expanduser(pattern))])

    return f"""#!/bin/bash
set -e
echo \"=== Pre-upgrade Script Started: $(date) ===\"

# Stop services
{stop_services}
systemctl daemon-reload

# Unload kernel modules
{unload_modules}

# Remove custom files and directories
{cleanup_cmds}

# Cleanup system
apt update
apt install -y --no-install-recommends bleachbit
bleachbit --clean system.cache system.tmp apt.autoclean apt.autoremove
apt purge -y bleachbit
apt autoremove -y
apt clean

# Backup module configs
mkdir -p ~/module_backup
cp /etc/modules-load.d/*stratopi*.conf ~/module_backup/
cp /etc/modprobe.d/*stratopi*.conf ~/module_backup/

echo \"=== Pre-upgrade Script Finished: $(date) ===\"
"""


def upgrade_script():
    return """#!/bin/bash
set -e
echo \"=== Upgrade Script Started: $(date) ===\"
apt update
apt full-upgrade -y
apt install raspberrypi-kernel-headers build-essential dkms -y
apt clean
echo \"=== Upgrade Script Finished: $(date) ===\"
"""


def post_upgrade_script(kernel_modules):
    reload_modules = "\n".join([f"modprobe {mod}" for mod in kernel_modules])

    return f"""#!/bin/bash
set -e
echo \"=== Post-upgrade Script Started: $(date) ===\"

# Recompile kernel modules
cd /usr/src/stratopi-modules
make clean
make
make install
depmod -a

# Reload kernel modules
{reload_modules}

# Verify modules loaded
dmesg | grep stratopi

apt clean
rm -rf /var/cache/apt/* /tmp/* /var/tmp/*

# Clean temporary files
rm -rf ~/module_backup

echo \"=== Post-upgrade Script Finished: $(date) ===\"
"""


def pack_logs():
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    archive_name = f"upgrade_logs_{timestamp}.tar.gz"
    subprocess.run(["tar", "czf", archive_name, str(LOG_DIR)], check=True)
    return archive_name


def send_logs(archive_name, config):
    scp_command = [
        "sshpass", "-p", config["ssh_pass"],
        "scp", "-P", str(config["remote_port"]),
        archive_name, f"{config['remote_user']}@{config['remote_server']}:{config['remote_path']}/"
    ]
    subprocess.run(scp_command, check=True)


def main():
    parser = argparse.ArgumentParser(description="Raspberry Pi Upgrade Orchestrator")
    parser.add_argument("-c", "--config", help="Path to config file (JSON format)")
    args = parser.parse_args()

    config = load_config(args.config)

    run_script(pre_upgrade_script(config["services"], config["kernel_modules"], config["cleanup_paths"]),
               "pre_upgrade.log")
    run_script(upgrade_script(), "upgrade.log")
    run_script(post_upgrade_script(config["kernel_modules"]), "post_upgrade.log")

    archive_name = pack_logs()
    send_logs(archive_name, config)

    subprocess.run(["sudo", "reboot"])


if __name__ == "__main__":
    main()