agent-orchestrator/agor-pro/src/branch_policy.rs

// SPDX-License-Identifier: LicenseRef-Commercial
// Branch Policy Enforcement — block agent sessions on protected branches.

use rusqlite::params;
use serde::Serialize;
use std::process::Command;

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct BranchPolicy {
    pub id: i64,
    pub pattern: String,
    pub action: String,
    pub reason: String,
}

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct PolicyDecision {
    pub allowed: bool,
    pub branch: String,
    pub matched_policy: Option<BranchPolicy>,
    pub reason: String,
}

fn ensure_tables(conn: &rusqlite::Connection) -> Result<(), String> {
    conn.execute_batch(
        "CREATE TABLE IF NOT EXISTS pro_branch_policies (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            pattern TEXT NOT NULL,
            action TEXT NOT NULL DEFAULT 'block',
            reason TEXT NOT NULL DEFAULT ''
        );"
    ).map_err(|e| format!("Failed to create branch_policies table: {e}"))?;

    // Seed default policies if table is empty
    let count: i64 = conn.query_row(
        "SELECT COUNT(*) AS cnt FROM pro_branch_policies", [], |row| row.get("cnt")
    ).unwrap_or(0);

    if count == 0 {
        conn.execute_batch(
            "INSERT INTO pro_branch_policies (pattern, action, reason) VALUES
                ('main', 'block', 'Protected branch: direct work on main is not allowed'),
                ('master', 'block', 'Protected branch: direct work on master is not allowed'),
                ('release/*', 'block', 'Protected branch: release branches require PRs');"
        ).map_err(|e| format!("Failed to seed default policies: {e}"))?;
    }

    Ok(())
}

/// Simple glob matching: supports `*` at the end of a pattern (e.g., `release/*`).
fn glob_match(pattern: &str, value: &str) -> bool {
    if pattern == value {
        return true;
    }
    if let Some(prefix) = pattern.strip_suffix('*') {
        return value.starts_with(prefix);
    }
    false
}

fn get_current_branch(project_path: &str) -> Result<String, String> {
    if project_path.starts_with('-') {
        return Err("Invalid project path: cannot start with '-'".into());
    }
    let output = Command::new("git")
        .args(["-C", project_path, "branch", "--show-current"])
        .output()
        .map_err(|e| format!("Failed to run git: {e}"))?;

    if output.status.success() {
        Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
    } else {
        Err("Not a git repository or git not available".into())
    }
}

#[tauri::command]
pub fn pro_branch_check(project_path: String) -> Result<PolicyDecision, String> {
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;

    let branch = get_current_branch(&project_path)?;

    let mut stmt = conn.prepare(
        "SELECT id, pattern, action, reason FROM pro_branch_policies"
    ).map_err(|e| format!("Query failed: {e}"))?;

    let policies: Vec<BranchPolicy> = stmt.query_map([], |row| {
        Ok(BranchPolicy {
            id: row.get("id")?,
            pattern: row.get("pattern")?,
            action: row.get("action")?,
            reason: row.get("reason")?,
        })
    }).map_err(|e| format!("Query failed: {e}"))?
    .collect::<Result<Vec<_>, _>>()
    .map_err(|e| format!("Row read failed: {e}"))?;

    for policy in &policies {
        if glob_match(&policy.pattern, &branch) {
            let allowed = policy.action != "block";
            return Ok(PolicyDecision {
                allowed,
                branch: branch.clone(),
                matched_policy: Some(policy.clone()),
                reason: policy.reason.clone(),
            });
        }
    }

    Ok(PolicyDecision {
        allowed: true,
        branch,
        matched_policy: None,
        reason: "No matching policy".into(),
    })
}

#[tauri::command]
pub fn pro_branch_policy_list() -> Result<Vec<BranchPolicy>, String> {
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;

    let mut stmt = conn.prepare(
        "SELECT id, pattern, action, reason FROM pro_branch_policies ORDER BY id"
    ).map_err(|e| format!("Query failed: {e}"))?;

    let rows = stmt.query_map([], |row| {
        Ok(BranchPolicy {
            id: row.get("id")?,
            pattern: row.get("pattern")?,
            action: row.get("action")?,
            reason: row.get("reason")?,
        })
    }).map_err(|e| format!("Query failed: {e}"))?
    .collect::<Result<Vec<_>, _>>()
    .map_err(|e| format!("Row read failed: {e}"))?;

    Ok(rows)
}

#[tauri::command]
pub fn pro_branch_policy_add(pattern: String, action: Option<String>, reason: Option<String>) -> Result<i64, String> {
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;
    let act = action.unwrap_or_else(|| "block".into());
    if !["block", "warn"].contains(&act.as_str()) {
        return Err(format!("Invalid action '{}': must be 'block' or 'warn'", act));
    }
    let rsn = reason.unwrap_or_default();
    conn.execute(
        "INSERT INTO pro_branch_policies (pattern, action, reason) VALUES (?1, ?2, ?3)",
        params![pattern, act, rsn],
    ).map_err(|e| format!("Failed to add policy: {e}"))?;
    Ok(conn.last_insert_rowid())
}

#[tauri::command]
pub fn pro_branch_policy_remove(id: i64) -> Result<(), String> {
    let conn = super::open_sessions_db()?;
    ensure_tables(&conn)?;
    conn.execute("DELETE FROM pro_branch_policies WHERE id = ?1", params![id])
        .map_err(|e| format!("Failed to remove policy: {e}"))?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_policy_decision_serializes_camel_case() {
        let d = PolicyDecision {
            allowed: false,
            branch: "main".into(),
            matched_policy: Some(BranchPolicy {
                id: 1,
                pattern: "main".into(),
                action: "block".into(),
                reason: "Protected".into(),
            }),
            reason: "Protected".into(),
        };
        let json = serde_json::to_string(&d).unwrap();
        assert!(json.contains("matchedPolicy"));
        assert!(json.contains("\"allowed\":false"));
    }

    #[test]
    fn test_branch_policy_serializes_camel_case() {
        let p = BranchPolicy {
            id: 1,
            pattern: "release/*".into(),
            action: "block".into(),
            reason: "No direct commits".into(),
        };
        let json = serde_json::to_string(&p).unwrap();
        assert!(json.contains("\"pattern\":\"release/*\""));
        assert!(json.contains("\"action\":\"block\""));
    }

    #[test]
    fn test_glob_match() {
        assert!(glob_match("main", "main"));
        assert!(!glob_match("main", "main2"));
        assert!(glob_match("release/*", "release/v1.0"));
        assert!(glob_match("release/*", "release/hotfix"));
        assert!(!glob_match("release/*", "feature/test"));
        assert!(!glob_match("master", "main"));
    }
}