#!/usr/bin/env python3
"""
Drop-in reference: /opt/sshportal/app.py
SSH-Portal mit Telegram-2FA und JIT-Firewall-Freischaltung.
Login mehlzoerwer -> Erstpasswort -> PW-Wechsel-Pflicht -> 2. Login: PW+Telegram-Code
-> UFW oeffnet Port 22 fuer angefragte IP (1 Stunde) -> Auto-Close.
Abhaengigkeiten (venv): flask, werkzeug, requests.
Start via systemd unit (NICHT nohup & im Agent-Shell).
"""
import os, re, time, secrets, sqlite3, subprocess, threading
from pathlib import Path
from flask import Flask, request, render_template_string, session, redirect, url_for
from werkzeug.security import generate_password_hash, check_password_hash
import requests

PORTAL_DIR = Path("/opt/sshportal")
DB = PORTAL_DIR / "portal.db"
LOG = PORTAL_DIR / "portal.log"
ALLOW_TIME = 3600          # 1 Stunde offen
MAX_2FA_ATTEMPTS = 3
CODE_TTL = 120             # 2 Minuten gueltig

ENV_FILE = "/root/.hermes/.env"
def _env(key):
    try:
        txt = Path(ENV_FILE).read_text()
        m = re.search(key + r"\s*=\s*(\S+)", txt)
        return m.group(1).strip() if m else ""
    except Exception:
        return ""
TG_TOKEN = _env("TELEGRAM_BOT_TOKEN")
TG_CHAT = _env("TELEGRAM_HOME_CHANNEL") or "6417835766"

app = Flask(__name__)
app.secret_key = secrets.token_hex(32)

def log(msg):
    with open(LOG, "a") as f:
        f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} {msg}\n")

def db():
    c = sqlite3.connect(DB); c.row_factory = sqlite3.Row; return c

def init_db():
    c = db()
    c.execute("""CREATE TABLE IF NOT EXISTS users(
        id INTEGER PRIMARY KEY, username TEXT UNIQUE, pw_hash TEXT,
        pw_changed INTEGER DEFAULT 0, failed INTEGER DEFAULT 0, lockout_until INTEGER DEFAULT 0)""")
    c.execute("""CREATE TABLE IF NOT EXISTS sessions(
        ip TEXT, status TEXT, ts INTEGER, expiry INTEGER)""")
    c.commit(); c.close()

def ensure_first_user():
    c = db()
    row = c.execute("SELECT * FROM users WHERE username='mehlzoerwer'").fetchone()
    if not row:
        c.execute("INSERT INTO users(username,pw_hash,pw_changed) VALUES(?,?,0)",
                  ("mehlzoerwer", generate_password_hash("St230918kM@vps")))
        c.commit(); log("Erstbenutzer angelegt (Erstpasswort aktiv)")
    c.close()

def send_telegram(text):
    try:
        requests.post(f"https://api.telegram.org/bot{TG_TOKEN}/sendMessage",
                      json={"chat_id": TG_CHAT, "text": text}, timeout=8)
    except Exception as e:
        log(f"Telegram-Fehler: {e}")

def ufw_allow(ip):
    r = subprocess.run(["ufw","status"], capture_output=True, text=True)
    if f"from {ip}" in r.stdout:
        return False
    subprocess.run(["ufw","allow","from",ip,"to","any","port","22","proto","tcp",
                    "comment","SSH JIT Portal"], check=False)
    return True

def ufw_deny(ip):
    subprocess.run(["ufw","delete","allow","from",ip,"to","any","port","22","proto","tcp",
                    "comment","SSH JIT Portal"], check=False)

def client_ip():
    return request.headers.get("X-Real-IP") or request.remote_addr

def _close_expired():
    while True:
        try:
            c = db(); now = int(time.time())
            for r in c.execute("SELECT ip FROM sessions WHERE status='open' AND expiry<=?", (now,)):
                ufw_deny(r["ip"])
                c.execute("UPDATE sessions SET status='closed' WHERE ip=? AND status='open'", (r["ip"],))
                log(f"Auto-Close: SSH geschlossen fuer {r['ip']}")
            c.commit(); c.close()
        except Exception:
            pass
        time.sleep(60)

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

HTML = """<!doctype html><html lang="de"><head><meta charset="utf-8"><title>SSH-Zugangsportal</title>
<meta name="viewport" content="width=device-width, initial-scale=1.0"></head>
<body style="font-family:system-ui;max-width:400px;margin:80px auto;padding:0 20px">
{{ message }}{{ content }}</body></html>"""

def render(content, msg=None, msgtype="info"):
    from markupsafe import Markup
    m = ""
    if msg:
        m = f'<div style="background:{ "#e6f9ef;color:#0a7a4a" if msgtype=="ok" else "#fdecea;color:#b3261e" };padding:10px;border-radius:5px;margin:10px 0">{msg}</div>'
    return render_template_string(HTML, message=Markup(m), content=Markup(content))

FORM_LOGIN = '<h1>SSH-Zugang</h1><form method="post">' + \
    '<p><input name="username" placeholder="Benutzer" required style="width:100%;padding:10px"></p>' + \
    '<p><input name="password" type="password" placeholder="Passwort" required style="width:100%;padding:10px"></p>' + \
    '<button style="width:100%;padding:11px;background:#10b981;color:#fff;border:0;border-radius:5px">Anmelden</button></form>'

FORM_PW = '<h1>Neues Passwort</h1><form method="post">' + \
    '<p><input name="new1" type="password" placeholder="Neues Passwort" required style="width:100%;padding:10px"></p>' + \
    '<p><input name="new2" type="password" placeholder="Wiederholen" required style="width:100%;padding:10px"></p>' + \
    '<button style="width:100%;padding:11px;background:#10b981;color:#fff;border:0;border-radius:5px">Speichern</button></form>'

FORM_2FA = '<h1>2-Faktor</h1><p>Ein 6-stelliger Code wurde dir per Telegram zugesendet. Gültig 2 Minuten.</p>' + \
    '<form method="post"><input name="code" placeholder="6-stelliger Code" required style="width:100%;padding:10px">' + \
    '<button style="width:100%;padding:11px;background:#10b981;color:#fff;border:0;border-radius:5px;margin-top:8px">Bestätigen</button></form>'

@app.route("/", methods=["GET","POST"])
def login():
    if request.method == "POST":
        username = request.form.get("username","").strip()
        password = request.form.get("password","")
        ip = client_ip()
        c = db()
        user = c.execute("SELECT * FROM users WHERE username=?", (username,)).fetchone()
        if not user or not check_password_hash(user["pw_hash"], password):
            log(f"FEHL_LOGIN {username} von {ip}")
            return render(FORM_LOGIN, "Falscher Benutzer oder Passwort.", "err")
        session["user"] = username
        if not user["pw_changed"]:
            log(f"PW-WECHSEL noetig {username} von {ip}")
            return redirect(url_for("change_pw"))
        session["pending_2fa_ip"] = ip
        code = f"{secrets.randbelow(1000000):06d}"
        session["2fa_code"] = code; session["2fa_ts"] = int(time.time()); session["2fa_attempts"] = 0
        send_telegram(f"SSH-Portal 2FA fuer Login von {ip}\nCode: {code}\nGueltig 2 Min.")
        log(f"2FA-Code an Chat {TG_CHAT} fuer Login von {ip}")
        return redirect(url_for("twofa"))
    return render(FORM_LOGIN)

@app.route("/change-pw", methods=["GET","POST"])
def change_pw():
    if "user" not in session:
        return redirect(url_for("login"))
    if request.method == "POST":
        new1 = request.form.get("new1",""); new2 = request.form.get("new2","")
        if new1 != new2 or len(new1) < 10:
            return render(FORM_PW, "Passwoerter mismatch oder zu kurz (min 10).", "err")
        c = db()
        c.execute("UPDATE users SET pw_hash=?, pw_changed=1 WHERE username=?",
                  (generate_password_hash(new1), session["user"]))
        c.commit(); c.close()
        log(f"Passwort geaendert fuer {session['user']}")
        send_telegram(f"Passwort fuer {session['user']} geaendert.")
        return render("<h1>Passwort geaendert</h1><p>Bitte neu einloggen. <a href='/'>Login</a></p>",
                      "Passwort erfolgreich geaendert.", "ok")
    return render(FORM_PW)

@app.route("/2fa", methods=["GET","POST"])
def twofa():
    if "2fa_code" not in session:
        return redirect(url_for("login"))
    if request.method == "POST":
        entered = request.form.get("code","").strip()
        ip = session["pending_2fa_ip"]
        if int(time.time()) - session["2fa_ts"] > CODE_TTL:
            session.pop("2fa_code", None)
            return render("", "Code abgelaufen. Bitte neu anmelden.", "err")
        if entered == session["2fa_code"]:
            added = ufw_allow(ip)
            c = db()
            c.execute("INSERT INTO sessions(ip,status,ts,expiry) VALUES(?,?,?,?)",
                      (ip,"open",int(time.time()),int(time.time())+ALLOW_TIME))
            c.commit(); c.close()
            session.pop("2fa_code", None)
            log(f"SSH-Freischaltung fuer {ip} (1h)")
            send_telegram(f"SSH-Freischaltung fuer {ip} aktiv (1 Stunde).")
            return render(f"<h1>Zugriff freigeschaltet</h1><p>SSH auf Port 22 fuer <b>{ip}</b> fuer 1 Stunde.</p>",
                          "Freigeschaltet!", "ok")
        else:
            session["2fa_attempts"] = session.get("2fa_attempts",0)+1
            log(f"FALSCHER 2FA-Code von {ip} (Versuch {session['2fa_attempts']})")
            if session["2fa_attempts"] >= MAX_2FA_ATTEMPTS:
                session.pop("2fa_code", None)
                send_telegram(f"3x falscher 2FA-Code von {ip} — gesperrt.")
                return render("", "Zu viele Fehlversuche. Login gesperrt.", "err")
            return render(FORM_2FA, f"Falscher Code. Versuche uebrig: {MAX_2FA_ATTEMPTS-session['2fa_attempts']}", "err")
    return render(FORM_2FA)

if __name__ == "__main__":
    PORTAL_DIR.mkdir(exist_ok=True)
    init_db(); ensure_first_user()
    log("SSH-Portal gestartet")
    app.run(host="127.0.0.1", port=8091)