main.py

# -*- coding: utf-8 -*-
import os
import shutil
import time
import threading
import subprocess
import urllib.request
from pathlib import Path

# Load Configs
SPACE_2_URL = os.environ.get("SPACE_2_URL", "https://augment17-claude-code-backend.hf.space")
BACKEND_API_KEY = os.environ.get("BACKEND_API_KEY", "")
BACKUP_GIT_REPO = os.environ.get("BACKUP_GIT_REPO", "") # SSH target e.g. git@github.com:augment17/my-backup.git
SSH_PRIVATE_KEY = os.environ.get("SSH_PRIVATE_KEY", "") # Secret private key string
BACKUP_LOCAL_DIR = "/tmp/git_backup_repo"
MAX_PART_SIZE_MB = 50
SYNC_INTERVAL_SECONDS = 300

def setup_ssh():
    """Setup SSH private key for GitHub authentications if provided."""
    if not SSH_PRIVATE_KEY:
        print("[SSH] No SSH_PRIVATE_KEY env variable found. Using default git authentication...")
        return
        
    ssh_dir = Path("/home/user/.ssh") if os.path.exists("/home/user") else Path("/tmp/.ssh")
    ssh_dir.mkdir(parents=True, exist_ok=True)
    
    key_file = ssh_dir / "id_rsa"
    key_file.write_text(SSH_PRIVATE_KEY.strip() + "\n")
    key_file.chmod(0o600)
    
    # Configure git to use this key and ignore strict host checks
    ssh_cmd = f"ssh -i {key_file} -o StrictHostKeyChecking=no"
    os.environ["GIT_SSH_COMMAND"] = ssh_cmd
    print("[SSH] Private key configured successfully.")

def run_cmd(cmd, cwd=None):
    try:
        subprocess.run(cmd, shell=True, check=True, cwd=cwd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        return True
    except subprocess.CalledProcessError:
        return False

def backup_loop():
    print("[Backup] Daemon started.")
    while True:
        try:
            # 1. Clean local backup directory
            if os.path.exists(BACKUP_LOCAL_DIR):
                shutil.rmtree(BACKUP_LOCAL_DIR)
            os.makedirs(BACKUP_LOCAL_DIR, exist_ok=True)

            # 2. Download zip from Space 2
            archive_path = Path("/tmp/workspace_backup.zip")
            if archive_path.exists():
                archive_path.unlink()

            download_url = f"{SPACE_2_URL}/api/backup/download"
            headers = {"Authorization": f"Bearer {BACKEND_API_KEY}"}
            
            print(f"[Backup] Downloading snapshot from {SPACE_2_URL}...")
            req = urllib.request.Request(download_url, headers=headers)
            with urllib.request.urlopen(req) as response:
                with open(archive_path, 'wb') as out_file:
                    shutil.copyfileobj(response, out_file)
            print("[Backup] Download completed successfully.")

            # 3. Handle Split if needed
            target_dest = Path(BACKUP_LOCAL_DIR) / "workspace_backup.zip"
            if archive_path.stat().st_size > (MAX_PART_SIZE_MB * 1024 * 1024):
                print(f"[Backup] File exceeds {MAX_PART_SIZE_MB}MB. Splitting...")
                run_cmd(f"split -b {MAX_PART_SIZE_MB}M {archive_path} {target_dest}.part")
            else:
                shutil.copy(archive_path, target_dest)

            # Clean up local zip
            archive_path.unlink()

            # 4. Git init and force push
            run_cmd("git init", cwd=BACKUP_LOCAL_DIR)
            run_cmd("git config user.name 'Backup Agent'", cwd=BACKUP_LOCAL_DIR)
            run_cmd("git config user.email 'backup@agent.internal'", cwd=BACKUP_LOCAL_DIR)
            run_cmd(f"git remote add origin {BACKUP_GIT_REPO}", cwd=BACKUP_LOCAL_DIR)
            run_cmd("git checkout -b main", cwd=BACKUP_LOCAL_DIR)
            run_cmd("git add -A", cwd=BACKUP_LOCAL_DIR)
            run_cmd('git commit -m "Auto-backup: ' + time.strftime("%Y-%m-%d %H:%M:%S") + '"', cwd=BACKUP_LOCAL_DIR)
            
            success = run_cmd("git push origin main --force", cwd=BACKUP_LOCAL_DIR)
            if success:
                print("[Backup] Sync completed successfully (git history purged).")
            else:
                print("[Backup Error] Git push failed.")
        except Exception as e:
            print(f"[Backup Error] Unexpected error: {e}")
        time.sleep(SYNC_INTERVAL_SECONDS)

def ping_loop():
    print("[Pinger] Daemon started.")
    while True:
        try:
            # Keep Space 2 awake
            health_url = f"{SPACE_2_URL}/health"
            req = urllib.request.Request(health_url)
            with urllib.request.urlopen(req) as response:
                status = response.status
                print(f"[Pinger] Pinged Space 2 health: Status {status}")
        except Exception as e:
            print(f"[Pinger Error] Failed to ping Space 2: {e}")
        time.sleep(SYNC_INTERVAL_SECONDS)

if __name__ == "__main__":
    setup_ssh()
    # Start loops in separate threads
    threading.Thread(target=ping_loop, daemon=True).start()
    backup_loop()