import gradio as gr
import subprocess
import os
import time
import threading
import glob
import sys

# Configuration
SSH_PORT = 2222
CF_BINARY = "./cloudflared"
LOG_FILE = "cloudflared.log"
SSHD_LOG = "sshd.log"
SSH_DIR = "/code/ssh"

def reassemble_binary():
    if os.path.exists(CF_BINARY):
        return True
    try:
        chunks = sorted(glob.glob("cf_part_*"))
        if not chunks: return False
        with open(CF_BINARY, "wb") as outfile:
            for chunk in chunks:
                with open(chunk, "rb") as infile:
                    outfile.write(infile.read())
        os.chmod(CF_BINARY, 0o777)
        return True
    except: return False

def generate_config():
    # HARDENED CONFIG FOR NON-ROOT + PERMISSIVE
    config = f"""
Port {SSH_PORT}
ListenAddress 127.0.0.1
PermitRootLogin yes
PasswordAuthentication yes
ChallengeResponseAuthentication no
UsePAM no
PidFile {SSH_DIR}/sshd.pid
HostKey {SSH_DIR}/ssh_host_rsa_key
AuthorizedKeysFile {SSH_DIR}/authorized_keys
Subsystem sftp /usr/lib/openssh/sftp-server
StrictModes no
UsePrivilegeSeparation no
"""
    with open(f"{SSH_DIR}/sshd_config", "w") as f:
        f.write(config)

def start_services():
    try:
        if not os.path.exists(SSH_DIR):
            os.makedirs(SSH_DIR, exist_ok=True)
            
        if not reassemble_binary():
            print("❌ Binary failed")
            return

        generate_config()
        
        # Keys
        key_path = os.path.join(SSH_DIR, "ssh_host_rsa_key")
        if not os.path.exists(key_path):
            subprocess.run(["ssh-keygen", "-t", "rsa", "-f", key_path, "-N", ""])
            
        # SSHD (Debug Mode to Log)
        with open(SSHD_LOG, "w") as f_out:
            print("🚀 Launching SSHD (Hardened)...")
            # -D: No detach, -e: Log to stderr, -d: Debug
            subprocess.Popen(["/usr/sbin/sshd", "-f", f"{SSH_DIR}/sshd_config", "-D", "-e"], stdout=f_out, stderr=f_out)
            
        # Cloudflared
        cmd = f"{CF_BINARY} tunnel --url ssh://localhost:{SSH_PORT} --logfile {LOG_FILE} --metrics localhost:45678"
        print(f"🚇 Launching Tunnel...")
        subprocess.Popen(cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            
    except Exception as e:
        print(f"❌ Error: {e}")

def get_logs():
    logs = "--- SSHD LOGS ---\n"
    if os.path.exists(SSHD_LOG):
        with open(SSHD_LOG, "r") as f: logs += f.read()[-1000:]
    logs += "\n\n--- CLOUDFLARED LOGS ---\n"
    if os.path.exists(LOG_FILE):
        with open(LOG_FILE, "r") as f:
            lines = f.readlines()
            for line in lines:
                if "trycloudflare.com" in line:
                    logs = f"🎯 FOUND: {line}\n" + logs
            logs += "".join(lines[-20:])
    return logs

threading.Thread(target=start_services, daemon=True).start()

with gr.Blocks() as demo:
    gr.Markdown("# 🚇 Brain13: Hardened SSH")
    log_box = gr.TextArea(label="Logs", lines=20, interactive=False)
    demo.load(get_logs, None, log_box, every=5)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)