action-gateway/src/routes.rs

use argon2::{Argon2, PasswordHash, PasswordVerifier};
use axum::{
    extract::{Form, Path, Query, State},
    http::{HeaderMap, StatusCode},
    response::{Html, IntoResponse, Redirect},
    routing::{get, post},
    Json, Router,
};
use axum_extra::extract::cookie::{Cookie, CookieJar};
use chrono::Utc;
use serde::Deserialize;
use std::process::Command;
use uuid::Uuid;

use crate::{models::*, AppState};

pub fn router(state: AppState) -> Router {
    Router::new()
        // Health check (no auth)
        .route("/health", get(health_check))
        // Web UI
        .route("/", get(dashboard))
        .route("/login", get(login_page).post(login_post))
        .route("/logout", get(logout))
        .route("/actions/:id/approve", post(approve_action))
        .route("/actions/:id/reject", post(reject_action))
        .route("/actions/:id", get(action_detail))
        // Agent API
        .route("/api/actions", post(submit_action))
        .route("/api/actions/:id/result", get(get_result))
        .with_state(state)
}

// ── Health Check ─────────────────────────────────────────────────

async fn health_check(State(state): State<AppState>) -> impl IntoResponse {
    // Verify the database is accessible with a simple query
    let db_ok = {
        let db = state.db.lock().unwrap();
        db.query_row("SELECT 1", [], |_| Ok(())).is_ok()
    };

    if db_ok {
        (
            StatusCode::OK,
            Json(serde_json::json!({"ok": true, "status": "live"})),
        )
            .into_response()
    } else {
        (
            StatusCode::SERVICE_UNAVAILABLE,
            Json(serde_json::json!({"ok": false, "status": "database unreachable"})),
        )
            .into_response()
    }
}

// ── Helpers ──────────────────────────────────────────────────────

fn session_user(jar: &CookieJar, state: &AppState) -> Option<String> {
    let token = jar.get("session")?.value().to_string();
    let db = state.db.lock().unwrap();
    db.query_row(
        "SELECT u.username FROM sessions s JOIN users u ON u.id = s.user_id 
         WHERE s.token = ?1 AND s.expires_at > datetime('now')",
        [&token],
        |row| row.get(0),
    )
    .ok()
}

fn verify_api_key(headers: &HeaderMap, state: &AppState) -> bool {
    let auth = headers
        .get("authorization")
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.strip_prefix("Bearer "))
        .unwrap_or("")
        .to_string();
    if auth.is_empty() {
        return false;
    }
    let db = state.db.lock().unwrap();
    let keys: Vec<String> = {
        let mut stmt = db.prepare("SELECT key_hash FROM api_keys").unwrap();
        stmt.query_map([], |row| row.get(0))
            .unwrap()
            .filter_map(|r| r.ok())
            .collect()
    };
    keys.iter().any(|hash| {
        let parsed = PasswordHash::new(hash).ok();
        parsed
            .map(|h| {
                Argon2::default()
                    .verify_password(auth.as_bytes(), &h)
                    .is_ok()
            })
            .unwrap_or(false)
    })
}

// ── Web UI ────────────────────────────────────────────────────────

async fn login_page() -> Html<String> {
    Html(page(
        "Login",
        r#"
    <div class="card" style="max-width:400px;margin:80px auto">
      <h2>🔐 Action Gateway</h2>
      <form method="post" action="/login">
        <label>Username</label>
        <input name="username" type="text" required autofocus>
        <label>Password</label>
        <input name="password" type="password" required>
        <button type="submit">Sign In</button>
      </form>
    </div>
    "#,
    ))
}

async fn login_post(
    State(state): State<AppState>,
    jar: CookieJar,
    Form(form): Form<LoginForm>,
) -> impl IntoResponse {
    let result: Option<(String, String)> = {
        let db = state.db.lock().unwrap();
        db.query_row(
            "SELECT id, password_hash FROM users WHERE username = ?1",
            [&form.username],
            |row| Ok((row.get(0)?, row.get(1)?)),
        )
        .ok()
    };

    if let Some((user_id, hash)) = result {
        let parsed = PasswordHash::new(&hash).unwrap();
        if Argon2::default()
            .verify_password(form.password.as_bytes(), &parsed)
            .is_ok()
        {
            let token = Uuid::new_v4().to_string();
            let expires = (Utc::now() + chrono::Duration::days(7))
                .format("%Y-%m-%d %H:%M:%S")
                .to_string();
            {
                let db = state.db.lock().unwrap();
                db.execute(
                    "INSERT INTO sessions (id, user_id, token, expires_at) VALUES (?1, ?2, ?3, ?4)",
                    [&Uuid::new_v4().to_string(), &user_id, &token, &expires],
                )
                .unwrap();
            }
            let cookie = Cookie::build(("session", token))
                .path("/")
                .http_only(true)
                .build();
            return (jar.add(cookie), Redirect::to("/")).into_response();
        }
    }
    Html(page(
        "Login",
        r#"
    <div class="card" style="max-width:400px;margin:80px auto">
      <h2>🔐 Action Gateway</h2>
      <p style="color:#e74c3c">Invalid credentials.</p>
      <form method="post" action="/login">
        <label>Username</label>
        <input name="username" type="text" required autofocus>
        <label>Password</label>
        <input name="password" type="password" required>
        <button type="submit">Sign In</button>
      </form>
    </div>
    "#,
    ))
    .into_response()
}

async fn logout(jar: CookieJar) -> impl IntoResponse {
    let jar = jar.remove(Cookie::from("session"));
    (jar, Redirect::to("/login"))
}

async fn dashboard(State(state): State<AppState>, jar: CookieJar) -> impl IntoResponse {
    let Some(username) = session_user(&jar, &state) else {
        return Redirect::to("/login").into_response();
    };

    let actions: Vec<Action> = {
        let db = state.db.lock().unwrap();
        let mut stmt = db.prepare(
            "SELECT id, chain_id, sequence, description, command, run_as, status, result, approve_token, submitted_at, resolved_at 
             FROM actions ORDER BY submitted_at DESC LIMIT 50"
        ).unwrap();
        stmt.query_map([], |row| {
            Ok(Action {
                id: row.get(0)?,
                chain_id: row.get(1)?,
                sequence: row.get(2)?,
                description: row.get(3)?,
                command: row.get(4)?,
                run_as: row.get(5)?,
                status: row.get(6)?,
                result: row.get(7)?,
                approve_token: row.get(8)?,
                submitted_at: row.get(9)?,
                resolved_at: row.get(10)?,
            })
        })
        .unwrap()
        .filter_map(|r| r.ok())
        .collect()
    };

    let rows: String = actions
        .iter()
        .map(|a| {
            let status_badge = match a.status.as_str() {
                "pending" => r#"<span class="badge pending">⏳ Pending</span>"#,
                "approved" => r#"<span class="badge approved">✅ Approved</span>"#,
                "executed" => r#"<span class="badge executed">✔ Executed</span>"#,
                "rejected" => r#"<span class="badge rejected">❌ Rejected</span>"#,
                "failed" => r#"<span class="badge failed">💥 Failed</span>"#,
                _ => r#"<span class="badge">Unknown</span>"#,
            };
            let actions_html = if a.status == "pending" {
                format!(
                    r#"
            <form method="post" action="/actions/{}/approve" style="display:inline">
              <button class="btn-approve" type="submit">✅ Approve</button>
            </form>
            <form method="post" action="/actions/{}/reject" style="display:inline">
              <button class="btn-reject" type="submit">❌ Reject</button>
            </form>
            "#,
                    a.id, a.id
                )
            } else {
                a.result
                    .clone()
                    .map(|r| format!("<pre class='result'>{}</pre>", html_escape(&r)))
                    .unwrap_or_default()
            };
            let run_as_badge = a
                .run_as
                .as_ref()
                .map(|u| format!("<span class='run-as'>run as: {}</span>", u))
                .unwrap_or_default();
            format!(
                r#"
        <div class="action-card">
          <div class="action-header">
            {} {} {}
            <span class="time">{}</span>
          </div>
          <div class="description">{}</div>
          <pre class="command">$ {}</pre>
          {}
        </div>
        "#,
                status_badge,
                a.chain_id
                    .as_ref()
                    .map(|c| format!("<span class='chain'>chain: {}</span>", &c[..8]))
                    .unwrap_or_default(),
                run_as_badge,
                a.submitted_at,
                html_escape(&a.description),
                html_escape(&a.command),
                actions_html
            )
        })
        .collect();

    let content = format!(
        r#"
    <div class="header">
      <h1>🤖 Action Gateway</h1>
      <span>Logged in as <strong>{}</strong> &nbsp; <a href="/logout">Logout</a></span>
    </div>
    <div class="actions-list" hx-get="/" hx-trigger="every 5s" hx-swap="outerHTML" hx-select=".actions-list">
      {}
    </div>
    "#,
        username,
        if rows.is_empty() {
            "<p style='text-align:center;opacity:0.5'>No actions yet.</p>".into()
        } else {
            rows
        }
    );

    Html(page("Action Gateway", &content)).into_response()
}

#[derive(Deserialize)]
struct ApproveQuery {
    token: Option<String>,
}

async fn action_detail(
    State(state): State<AppState>,
    jar: CookieJar,
    Path(id): Path<String>,
    Query(q): Query<ApproveQuery>,
) -> impl IntoResponse {
    // Allow approve via token link (no login required)
    if let Some(token) = q.token {
        let action: Option<Action> = {
            let db = state.db.lock().unwrap();
            db.query_row(
                "SELECT id, chain_id, sequence, description, command, run_as, status, result, approve_token, submitted_at, resolved_at FROM actions WHERE id = ?1 AND approve_token = ?2",
                [&id, &token],
                |row| Ok(Action {
                    id: row.get(0)?, chain_id: row.get(1)?, sequence: row.get(2)?,
                    description: row.get(3)?, command: row.get(4)?, run_as: row.get(5)?,
                    status: row.get(6)?, result: row.get(7)?, approve_token: row.get(8)?,
                    submitted_at: row.get(9)?, resolved_at: row.get(10)?,
                }),
            ).ok()
        };
        if let Some(a) = action {
            if a.status == "pending" {
                let content = format!(
                    r#"
                <div class="card" style="max-width:600px;margin:40px auto">
                  <h2>Review Action</h2>
                  <p>{}</p>
                  <pre class="command">$ {}</pre>
                  <form method="post" action="/actions/{}/approve" style="display:inline">
                    <button class="btn-approve">✅ Approve</button>
                  </form>
                  <form method="post" action="/actions/{}/reject" style="display:inline">
                    <button class="btn-reject">❌ Reject</button>
                  </form>
                </div>
                "#,
                    html_escape(&a.description),
                    html_escape(&a.command),
                    a.id,
                    a.id
                );
                return Html(page("Approve Action", &content)).into_response();
            }
        }
        return Html(page("Invalid", "<p>Invalid or expired token.</p>")).into_response();
    }

    // Otherwise require login
    if session_user(&jar, &state).is_none() {
        return Redirect::to("/login").into_response();
    }
    Redirect::to("/").into_response()
}

async fn approve_action(
    State(state): State<AppState>,
    jar: CookieJar,
    Path(id): Path<String>,
) -> impl IntoResponse {
    if session_user(&jar, &state).is_none() {
        return Redirect::to("/login").into_response();
    }

    let action: Option<Action> = {
        let db = state.db.lock().unwrap();
        db.query_row(
            "SELECT id, chain_id, sequence, description, command, run_as, status, result, approve_token, submitted_at, resolved_at FROM actions WHERE id = ?1",
            [&id],
            |row| Ok(Action {
                id: row.get(0)?, chain_id: row.get(1)?, sequence: row.get(2)?,
                description: row.get(3)?, command: row.get(4)?, run_as: row.get(5)?,
                status: row.get(6)?, result: row.get(7)?, approve_token: row.get(8)?,
                submitted_at: row.get(9)?, resolved_at: row.get(10)?,
            }),
        ).ok()
    };

    if let Some(a) = action {
        if a.status == "pending" {
            // Build command with run_as support
            let output = match a.run_as.as_deref() {
                Some("root") => Command::new("sudo")
                    .args(["-n", "sh", "-c", &a.command])
                    .output(),
                Some(user) => Command::new("sudo")
                    .args(["-n", "-u", user, "sh", "-c", &a.command])
                    .output(),
                None => Command::new("sh").args(["-c", &a.command]).output(),
            };

            let (status, result) = match output {
                Ok(out) => {
                    let stdout = sanitize_output(&String::from_utf8_lossy(&out.stdout));
                    let stderr = sanitize_output(&String::from_utf8_lossy(&out.stderr));
                    let combined = if stderr.is_empty() {
                        stdout
                    } else {
                        format!("{}\nSTDERR:\n{}", stdout, stderr)
                    };
                    if out.status.success() {
                        ("executed", combined)
                    } else {
                        ("failed", combined)
                    }
                }
                Err(e) => ("failed", e.to_string()),
            };

            let db = state.db.lock().unwrap();
            db.execute(
                "UPDATE actions SET status = ?1, result = ?2, resolved_at = datetime('now') WHERE id = ?3",
                [status, &result, &a.id],
            ).unwrap();
        }
    }

    Redirect::to("/").into_response()
}

async fn reject_action(
    State(state): State<AppState>,
    jar: CookieJar,
    Path(id): Path<String>,
) -> impl IntoResponse {
    if session_user(&jar, &state).is_none() {
        return Redirect::to("/login").into_response();
    }

    // Reject this action and all pending actions in same chain
    let chain_id: Option<String> = {
        let db = state.db.lock().unwrap();
        db.query_row("SELECT chain_id FROM actions WHERE id = ?1", [&id], |row| {
            row.get(0)
        })
        .ok()
        .flatten()
    };

    let db = state.db.lock().unwrap();
    db.execute(
        "UPDATE actions SET status = 'rejected', resolved_at = datetime('now') WHERE id = ?1",
        [&id],
    )
    .unwrap();

    if let Some(cid) = chain_id {
        db.execute(
            "UPDATE actions SET status = 'rejected', resolved_at = datetime('now') WHERE chain_id = ?1 AND status = 'pending'",
            [&cid],
        ).unwrap();
    }

    Redirect::to("/").into_response()
}

// ── Agent API ─────────────────────────────────────────────────────

async fn submit_action(
    State(state): State<AppState>,
    headers: HeaderMap,
    Json(body): Json<SubmitAction>,
) -> impl IntoResponse {
    if !verify_api_key(&headers, &state) {
        return (
            StatusCode::UNAUTHORIZED,
            Json(serde_json::json!({"error": "unauthorized"})),
        )
            .into_response();
    }

    let id = Uuid::new_v4().to_string();
    let approve_token = Uuid::new_v4().to_string();
    let chain_id = body.chain_id.unwrap_or_else(|| Uuid::new_v4().to_string());
    let sequence = body.sequence.unwrap_or(0);

    {
        let db = state.db.lock().unwrap();
        db.execute(
            "INSERT INTO actions (id, chain_id, sequence, description, command, run_as, approve_token) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
            rusqlite::params![&id, &chain_id, sequence, &body.description, &body.command, &body.run_as, &approve_token],
        ).unwrap();
    }

    // Send notification webhook
    let approve_url = format!(
        "{}/actions/{}?token={}",
        state.config.base_url, id, approve_token
    );
    if let Some(webhook) = &state.config.notification_webhook {
        let webhook = webhook.clone();
        let token = state.config.notification_webhook_token.clone();
        let description = body.description.clone();
        let command = body.command.clone();
        let url = approve_url.clone();
        tokio::spawn(async move {
            let client = reqwest::Client::new();
            let mut req = client.post(&webhook)
                .json(&serde_json::json!({
                    "text": format!("[ActionGateway] New action needs approval.\nDescription: {}\nCommand: {}\nApprove: {}", description, command, url),
                    "mode": "now"
                }));
            if let Some(t) = token {
                req = req.header("Authorization", format!("Bearer {}", t));
            }
            let _ = req.send().await;
        });
    }

    (
        StatusCode::CREATED,
        Json(serde_json::json!({
            "id": id,
            "approve_url": approve_url,
            "status": "pending"
        })),
    )
        .into_response()
}

async fn get_result(
    State(state): State<AppState>,
    headers: HeaderMap,
    Path(id): Path<String>,
) -> impl IntoResponse {
    if !verify_api_key(&headers, &state) {
        return (
            StatusCode::UNAUTHORIZED,
            Json(serde_json::json!({"error": "unauthorized"})),
        )
            .into_response();
    }
    let db = state.db.lock().unwrap();
    let result: Option<(String, Option<String>)> = db
        .query_row(
            "SELECT status, result FROM actions WHERE id = ?1",
            [&id],
            |row| Ok((row.get(0)?, row.get(1)?)),
        )
        .ok();

    match result {
        Some((status, result)) => {
            // Build JSON manually: serde_json's json! macro does not reliably escape
            // newlines/control chars in Option<String> in this build environment.
            let result_json = match result {
                Some(s) => {
                    let escaped: String = s
                        .chars()
                        .flat_map(|c| match c {
                            '"' => vec!['\\', '"'],
                            '\\' => vec!['\\', '\\'],
                            '\n' => vec!['\\', 'n'],
                            '\r' => vec!['\\', 'r'],
                            '\t' => vec!['\\', 't'],
                            c if (c as u32) < 0x20 => {
                                format!("\\u{:04x}", c as u32).chars().collect()
                            }
                            c => vec![c],
                        })
                        .collect();
                    format!("\"{}\"", escaped)
                }
                None => "null".to_string(),
            };
            let body = format!(r#"{{"status":"{}","result":{}}}"#, status, result_json);
            (
                StatusCode::OK,
                [(axum::http::header::CONTENT_TYPE, "application/json")],
                body,
            )
                .into_response()
        }
        None => (
            StatusCode::NOT_FOUND,
            Json(serde_json::json!({"error": "not found"})),
        )
            .into_response(),
    }
}

// ── HTML helpers ──────────────────────────────────────────────────

fn html_escape(s: &str) -> String {
    s.replace('&', "&amp;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
        .replace('"', "&quot;")
}

/// Strip ANSI/VT100 escape sequences and other control characters from command output.
/// Keeps printable text, newlines, and tabs. This ensures stored results are clean
/// UTF-8 text that serializes safely to JSON and renders correctly in the UI.
fn sanitize_output(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    let mut chars = s.chars().peekable();
    while let Some(c) = chars.next() {
        match c {
            // ESC: swallow escape sequences
            '\x1b' => {
                match chars.peek().copied() {
                    // CSI sequence: ESC [ ... <letter>
                    Some('[') => {
                        chars.next();
                        for nc in chars.by_ref() {
                            if nc.is_ascii_alphabetic() {
                                break;
                            }
                        }
                    }
                    // Character set designations: ESC ( X  or  ESC ) X
                    Some('(' | ')') => {
                        chars.next();
                        chars.next();
                    }
                    // Any other ESC sequence: skip one char
                    _ => {
                        chars.next();
                    }
                }
            }
            // CR: collapse \r\n → \n, standalone \r → \n
            '\r' => {
                if chars.peek() != Some(&'\n') {
                    out.push('\n');
                }
            }
            // Keep newlines and tabs as-is
            '\n' | '\t' => out.push(c),
            // Drop other control characters (0x00–0x1f, 0x7f)
            c if c.is_control() => {}
            // Keep everything else
            c => out.push(c),
        }
    }
    out
}

fn page(title: &str, content: &str) -> String {
    format!(
        r#"<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>{}</title>
<script src="https://unpkg.com/htmx.org@1.9.10"></script>
<style>
  *, *::before, *::after {{ box-sizing: border-box; margin: 0; padding: 0; }}
  body {{ font-family: system-ui, sans-serif; background: #0f1117; color: #e0e0e0; min-height: 100vh; padding: 20px; }}
  a {{ color: #7eb8f7; }}
  h1, h2 {{ font-weight: 600; }}
  .header {{ display: flex; justify-content: space-between; align-items: center; padding: 0 0 24px 0; border-bottom: 1px solid #2a2d3a; margin-bottom: 24px; }}
  .card {{ background: #1a1d27; border: 1px solid #2a2d3a; border-radius: 10px; padding: 24px; }}
  .card h2 {{ margin-bottom: 16px; }}
  label {{ display: block; margin: 12px 0 4px; font-size: 0.9em; color: #aaa; }}
  input {{ width: 100%; padding: 10px 12px; background: #0f1117; border: 1px solid #2a2d3a; border-radius: 6px; color: #e0e0e0; font-size: 1em; }}
  input:focus {{ outline: none; border-color: #7eb8f7; }}
  button {{ padding: 10px 20px; border: none; border-radius: 6px; cursor: pointer; font-size: 0.95em; font-weight: 500; margin-top: 16px; }}
  .btn-approve {{ background: #27ae60; color: white; margin-right: 8px; }}
  .btn-approve:hover {{ background: #2ecc71; }}
  .btn-reject {{ background: #c0392b; color: white; }}
  .btn-reject:hover {{ background: #e74c3c; }}
  button[type="submit"] {{ background: #7eb8f7; color: #0f1117; width: 100%; margin-top: 16px; }}
  .actions-list {{ display: flex; flex-direction: column; gap: 12px; }}
  .action-card {{ background: #1a1d27; border: 1px solid #2a2d3a; border-radius: 10px; padding: 20px; }}
  .action-header {{ display: flex; align-items: center; gap: 10px; margin-bottom: 10px; flex-wrap: wrap; }}
  .badge {{ padding: 3px 10px; border-radius: 20px; font-size: 0.8em; font-weight: 600; }}
  .badge.pending  {{ background: #7f6500; color: #ffd700; }}
  .badge.approved {{ background: #1a5c33; color: #2ecc71; }}
  .badge.executed {{ background: #1a4060; color: #7eb8f7; }}
  .badge.rejected {{ background: #5c1a1a; color: #e74c3c; }}
  .badge.failed   {{ background: #5c1a1a; color: #ff6b6b; }}
  .chain {{ font-size: 0.75em; color: #888; font-family: monospace; }}
  .run-as {{ font-size: 0.75em; color: #b39ddb; font-family: monospace; background: #2a2040; padding: 2px 8px; border-radius: 10px; }}
  .time {{ font-size: 0.8em; color: #666; margin-left: auto; }}
  .description {{ margin-bottom: 10px; color: #ccc; }}
  .command {{ background: #0a0c13; border: 1px solid #2a2d3a; border-radius: 6px; padding: 12px; font-size: 0.9em; color: #a8d8a8; overflow-x: auto; margin-bottom: 12px; }}
  .result {{ background: #0a0c13; border: 1px solid #2a2d3a; border-radius: 6px; padding: 12px; font-size: 0.85em; color: #ccc; overflow-x: auto; max-height: 200px; overflow-y: auto; white-space: pre-wrap; word-break: break-all; }}
  pre {{ white-space: pre-wrap; }}
</style>
</head>
<body>
{}
</body>
</html>"#,
        title, content
    )
}