diff --git a/src/commands/bot.rs b/src/commands/bot.rs index f7731b0..6cc230a 100644 --- a/src/commands/bot.rs +++ b/src/commands/bot.rs @@ -1,7 +1,7 @@ use anyhow::{Context, Result}; use serde::{Deserialize, Serialize}; use serde_json::Value; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::fs; use std::path::PathBuf; use std::process::Stdio; @@ -27,7 +27,7 @@ struct ClaudeSession { impl ClaudeSession { /// Spawn a persistent claude process with stream-json I/O - async fn spawn() -> Result { + async fn spawn(user_mode: bool) -> Result { // Run claude inside a dedicated directory under config let work_dir = dirs::config_dir() .context("Could not find config directory")? @@ -39,15 +39,50 @@ impl ClaudeSession { let rules_path = work_dir.join("CLAUDE.md"); fs::write(&rules_path, BOT_RULES)?; - eprintln!("bot: claude working directory = {}", work_dir.display()); + // Manage settings.json: create in user mode, remove otherwise + let settings_dir = work_dir.join(".claude"); + let settings_path = settings_dir.join("settings.json"); + if user_mode { + fs::create_dir_all(&settings_dir)?; + let settings = serde_json::json!({ + "permissions": { + "deny": [ + "Bash", + "Edit(.claude/**)", + "Write(.claude/**)", + "Edit(.mcp.json)", + "Write(.mcp.json)" + ], + "allow": [ + "Read", + "Glob", + "Grep" + ] + } + }); + fs::write(&settings_path, serde_json::to_string_pretty(&settings)?)?; + } else { + // Remove restrictive settings from previous --user run + let _ = fs::remove_file(&settings_path); + } - let mut child = tokio::process::Command::new("claude") - .arg("--input-format") + eprintln!("bot: claude working directory = {}", work_dir.display()); + eprintln!("bot: user_mode = {}", user_mode); + + let mut cmd = tokio::process::Command::new("claude"); + cmd.arg("--input-format") .arg("stream-json") .arg("--output-format") .arg("stream-json") - .arg("--verbose") - .arg("--dangerously-skip-permissions") + .arg("--verbose"); + + if user_mode { + cmd.arg("--permission-mode").arg("dontAsk"); + } else { + cmd.arg("--dangerously-skip-permissions"); + } + + let mut child = cmd .current_dir(&work_dir) .stdin(Stdio::piped()) .stdout(Stdio::piped()) @@ -187,6 +222,9 @@ struct BotState { /// Timestamp of the last seen notification #[serde(default)] last_seen: Option, + /// Per-user daily message counts: DID -> (date "YYYY-MM-DD", count) + #[serde(default)] + user_counts: HashMap, } /// Parsed notification for processing @@ -242,6 +280,33 @@ fn load_admin_did() -> Result { .context("config.json missing 'did' field") } +/// Load daily rate limit from config.json (bot.limit, default 3) +fn load_user_limit() -> u32 { + token::load_config() + .ok() + .and_then(|c| c["bot"]["limit"].as_u64()) + .unwrap_or(3) as u32 +} + +/// Check if a user is within daily rate limit. Returns true if allowed. +/// Increments count on success. +fn check_rate_limit(state: &mut BotState, did: &str, limit: u32) -> bool { + let today = chrono::Utc::now().format("%Y-%m-%d").to_string(); + let entry = state.user_counts.entry(did.to_string()).or_insert_with(|| (today.clone(), 0)); + + // Reset if new day + if entry.0 != today { + *entry = (today, 0); + } + + if entry.1 >= limit { + false + } else { + entry.1 += 1; + true + } +} + /// Fetch notifications using the bot session async fn fetch_notifications( client: &XrpcClient, @@ -340,20 +405,22 @@ async fn post_reply( } /// Main bot entry point -pub async fn start(interval_secs: u64) -> Result<()> { +pub async fn start(interval_secs: u64, user_mode: bool) -> Result<()> { let admin_did = load_admin_did()?; + let user_limit = if user_mode { load_user_limit() } else { 0 }; eprintln!("bot: admin DID = {}", admin_did); + eprintln!("bot: user_mode = {}, limit = {}/day", user_mode, user_limit); eprintln!("bot: polling interval = {}s", interval_secs); let mut state = load_state(); // Spawn persistent Claude session - let mut session = ClaudeSession::spawn().await?; + let mut session = ClaudeSession::spawn(user_mode).await?; eprintln!("bot: claude session started"); eprintln!("bot: starting notification loop..."); loop { - if let Err(e) = poll_once(&admin_did, &mut state, &mut session).await { + if let Err(e) = poll_once(&admin_did, &mut state, &mut session, user_mode, user_limit).await { eprintln!("bot: poll error: {}", e); } tokio::time::sleep(std::time::Duration::from_secs(interval_secs)).await; @@ -365,6 +432,8 @@ async fn poll_once( admin_did: &str, state: &mut BotState, claude: &mut ClaudeSession, + user_mode: bool, + user_limit: u32, ) -> Result<()> { // Refresh bot session let session = auth::refresh_bot_session().await?; @@ -403,12 +472,22 @@ async fn poll_once( } } - // Admin filter: only respond to admin let author_did = match notif["author"]["did"].as_str() { Some(d) => d, None => continue, }; - if author_did != admin_did { + + if author_did == admin_did { + // Admin: always pass + } else if user_mode { + // User mode: check rate limit + if !check_rate_limit(state, author_did, user_limit) { + eprintln!("bot: rate limited user {} (limit {}/day)", author_did, user_limit); + state.processed.insert(uri.to_string()); + continue; + } + } else { + // Admin-only mode: skip non-admin continue; } @@ -437,7 +516,7 @@ async fn poll_once( Err(e) => { eprintln!("bot: claude error: {}, respawning session...", e); // Try to respawn session - match ClaudeSession::spawn().await { + match ClaudeSession::spawn(user_mode).await { Ok(new_session) => { *claude = new_session; eprintln!("bot: claude session respawned"); @@ -584,7 +663,7 @@ pub async fn start_chat(interval_secs: u64) -> Result<()> { let mut state = load_chat_state(); // Spawn persistent Claude session - let mut claude = ClaudeSession::spawn().await?; + let mut claude = ClaudeSession::spawn(false).await?; eprintln!("chat-bot: claude session started"); eprintln!("chat-bot: starting chat poll loop..."); @@ -692,7 +771,7 @@ async fn chat_poll_once( Ok(r) => r, Err(e) => { eprintln!("chat-bot: claude error: {}, respawning...", e); - match ClaudeSession::spawn().await { + match ClaudeSession::spawn(false).await { Ok(new_session) => { *claude = new_session; match claude.send(text).await { diff --git a/src/main.rs b/src/main.rs index decc2af..db1c34e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -287,6 +287,9 @@ enum BotCommands { /// Poll interval in seconds #[arg(short, long, default_value = "30")] interval: u64, + /// Allow all users (with daily rate limit from config bot.limit) + #[arg(long)] + user: bool, }, /// Start the DM chat bot (poll chat messages and reply) Chat { @@ -412,8 +415,8 @@ async fn main() -> Result<()> { } Commands::Bot { command } => { match command { - BotCommands::Start { interval } => { - commands::bot::start(interval).await?; + BotCommands::Start { interval, user } => { + commands::bot::start(interval, user).await?; } BotCommands::Chat { interval } => { commands::bot::start_chat(interval).await?;