Skip to main content

relintio_agent/
lib.rs

1pub mod obsidian;
2pub mod utils;
3pub mod middleware;
4
5use std::collections::HashMap;
6use std::env;
7use std::fs;
8use std::path::PathBuf;
9use std::sync::atomic::{AtomicBool, Ordering};
10use std::sync::{Arc, Mutex, RwLock};
11use std::time::{Duration, SystemTime, UNIX_EPOCH};
12
13use hmac::{Hmac, Mac};
14use sha2::Sha256;
15use serde::{Deserialize, Serialize};
16use serde_json::Value;
17
18type HmacSha256 = Hmac<Sha256>;
19
20const AGENT_VERSION: &str = "0.1.2";
21const THRESHOLDS: &[(&str, u32)] = &[
22    ("ALLOW", 0),
23    ("SLOW", 40),
24    ("CHALLENGE", 60),
25    ("DECOY", 75),
26    ("BLOCK", 85),
27];
28
29const DECOY_HTML: &str = r#"<!DOCTYPE html><html><head><meta charset="utf-8"><title>Scheduled Maintenance</title><style>body{background:#0a0a0a;color:#aaa;font-family:system-ui;display:flex;justify-content:center;align-items:center;min-height:100vh;margin:0}.box{text-align:center;max-width:480px}h1{font-size:1.5rem;color:#fff;margin:0 0 1rem}p{color:#666;font-size:.875rem}</style></head><body><div class="box"><h1>Scheduled Maintenance</h1><p>We are currently performing scheduled maintenance. Please try again later.</p><p style="color:#444;font-size:.75rem;margin-top:2rem">ETA: ~15 minutes</p></div></body></html>"#;
30
31const BLOCK_HTML: &str = r#"<!DOCTYPE html><html><head><meta charset="utf-8"><title>Access Denied</title><style>body{background:#050507;color:#fff;font-family:system-ui;display:flex;align-items:center;justify-content:center;height:100vh;margin:0}.card{background:#0E0E10;padding:40px;border-radius:20px;border:1px solid rgba(255,255,255,.1);text-align:center}h1{color:#ef4444;margin:0 0 1rem}</style></head><body><div class="card"><h1>Access Blocked</h1><p>Security policies have flagged this request as suspicious.</p></div></body></html>"#;
32
33#[derive(Debug, Clone, PartialEq, Eq)]
34pub enum Decision {
35    Allow,
36    Slow,
37    Challenge { redirect_url: String },
38    Decoy,
39    Block,
40}
41
42#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct RelintioConfig {
44    pub license_key: String,
45    pub api_url: String,
46    pub sync_interval_seconds: u64,
47}
48
49struct TokenBucket {
50    tokens: f64,
51    last_ts: f64,
52}
53
54pub struct RelintioAgent {
55    config: RelintioConfig,
56    pub(crate) rules: Arc<RwLock<Option<Value>>>,
57    synced_at: Arc<Mutex<u64>>,
58    next_sync_at: Arc<Mutex<u64>>,
59    sync_failures: Arc<Mutex<u8>>,
60    sync_in_progress: AtomicBool,
61    buckets: Arc<Mutex<HashMap<String, TokenBucket>>>,
62    client: reqwest::Client,
63    cache_path: PathBuf,
64}
65
66impl RelintioAgent {
67    pub fn new(config: RelintioConfig) -> Self {
68        let license_hash = format!("{:x}", md5::compute(config.license_key.as_bytes()));
69        let mut cache_dir = env::temp_dir();
70        cache_dir.push("relintio");
71        let _ = fs::create_dir_all(&cache_dir);
72        cache_dir.push(format!("up_rules_{}.json", &license_hash[..16]));
73
74        let agent = Self {
75            config,
76            rules: Arc::new(RwLock::new(None)),
77            synced_at: Arc::new(Mutex::new(0)),
78            next_sync_at: Arc::new(Mutex::new(0)),
79            sync_failures: Arc::new(Mutex::new(0)),
80            sync_in_progress: AtomicBool::new(false),
81            buckets: Arc::new(Mutex::new(HashMap::new())),
82            client: reqwest::Client::builder()
83                .connect_timeout(Duration::from_secs(5))
84                .timeout(Duration::from_secs(10))
85                .build()
86                .unwrap_or_else(|_| reqwest::Client::new()),
87            cache_path: cache_dir,
88        };
89
90        // Try load cache on startup
91        let _ = agent.load_cache_from_disk();
92        agent
93    }
94
95    /// Verifies the query string challenge token using HMAC-SHA256 signature checks.
96    pub fn verify_up_token(&self, token: &str) -> bool {
97        use base64::{engine::general_purpose, Engine as _};
98        let decoded_bytes = match general_purpose::STANDARD.decode(token) {
99            Ok(b) => b,
100            Err(_) => return false,
101        };
102        let decoded = match String::from_utf8(decoded_bytes) {
103            Ok(s) => s,
104            Err(_) => return false,
105        };
106
107        let parts: Vec<&str> = decoded.split("::").collect();
108        if parts.len() != 2 {
109            return false;
110        }
111
112        let ts_raw = parts[0];
113        let sig = parts[1];
114
115        let ts: u64 = match ts_raw.parse() {
116            Ok(val) => val,
117            Err(_) => return false,
118        };
119
120        let now = match SystemTime::now().duration_since(UNIX_EPOCH) {
121            Ok(duration) => duration.as_secs(),
122            Err(_) => return false,
123        };
124
125        if now.abs_diff(ts) > 120 {
126            return false;
127        }
128
129        let message = format!("{}|{}", ts, self.config.license_key);
130        let mut mac = match HmacSha256::new_from_slice(self.config.license_key.as_bytes()) {
131            Ok(m) => m,
132            Err(_) => return false,
133        };
134        mac.update(message.as_bytes());
135        let result = mac.finalize();
136        let calc_sig = hex::encode(result.into_bytes());
137
138        // Constant time comparison
139        let calc_bytes = calc_sig.as_bytes();
140        let sig_bytes = sig.as_bytes();
141        if calc_bytes.len() != sig_bytes.len() {
142            return false;
143        }
144
145        let mut diff = 0;
146        for i in 0..calc_bytes.len() {
147            diff |= calc_bytes[i] ^ sig_bytes[i];
148        }
149        diff == 0
150    }
151
152    /// Calculate the SHA256 hashed signature value for setting the passport cookie.
153    pub fn passport_value(&self) -> String {
154        use sha2::Digest;
155        let mut hasher = Sha256::new();
156        hasher.update(format!("verified{}", self.config.license_key));
157        hex::encode(hasher.finalize())
158    }
159
160    fn load_cache_from_disk(&self) -> Result<(), Box<dyn std::error::Error>> {
161        if self.cache_path.exists() {
162            let data = fs::read_to_string(&self.cache_path)?;
163            let parsed: Value = serde_json::from_str(&data)?;
164            let mut w = self.rules.write().unwrap_or_else(|poisoned| poisoned.into_inner());
165            *w = Some(parsed);
166        }
167        Ok(())
168    }
169
170    fn save_cache_to_disk(&self, val: &Value) -> Result<(), Box<dyn std::error::Error>> {
171        let serialized = serde_json::to_string(val)?;
172        fs::write(&self.cache_path, serialized)?;
173        Ok(())
174    }
175
176    /// Triggers rules fetch from the API and saves it. Async-safe.
177    pub async fn refresh_rules(&self, domain: &str) -> Result<(), Box<dyn std::error::Error>> {
178        let url = format!("{}/agent/verify", self.config.api_url.trim_end_matches('/'));
179        let body = serde_json::json!({
180            "license_key": self.config.license_key,
181            "domain": domain,
182            "protocol_version": 1,
183            "agent_kind": "rust",
184            "agent_version": AGENT_VERSION,
185            "capabilities": ["custom_rules", "telemetry"]
186        });
187        let res = self.client.post(&url)
188            .json(&body)
189            .send()
190            .await?
191            .error_for_status()?;
192
193        let val: Value = res.json().await?;
194        {
195            let mut w = self.rules.write().unwrap_or_else(|poisoned| poisoned.into_inner());
196            *w = Some(val.clone());
197        }
198        let mut sync = self.synced_at.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
199        *sync = unix_now();
200        let _ = self.save_cache_to_disk(&val);
201        Ok(())
202    }
203
204    /// Heartbeat signal trigger.
205    pub async fn send_heartbeat(&self, domain: &str) {
206        let url = format!("{}/agent/heartbeat", self.config.api_url.trim_end_matches('/'));
207        let body = serde_json::json!({
208            "license_key": self.config.license_key,
209            "domain": domain,
210            "agent_version": AGENT_VERSION,
211            "agent_kind": "rust",
212            "timestamp": unix_now()
213        });
214
215        let _ = self.client.post(&url)
216            .json(&body)
217            .timeout(Duration::from_secs(2))
218            .send()
219            .await;
220    }
221
222    async fn challenge_url(&self, domain: &str, path: &str) -> Option<String> {
223        use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
224
225        let endpoint = format!("{}/agent/challenge/init", self.config.api_url.trim_end_matches('/'));
226        let scheme = if domain == "localhost" || domain.starts_with("localhost:") || domain == "127.0.0.1" || domain.starts_with("127.0.0.1:") {
227            "http"
228        } else {
229            "https"
230        };
231        let return_url = format!("{}://{}{}", scheme, domain, path);
232        let response = self.client.post(endpoint)
233            .json(&serde_json::json!({
234                "license_key": self.config.license_key,
235                "return_url": return_url,
236            }))
237            .send()
238            .await
239            .ok()?
240            .error_for_status()
241            .ok()?;
242        let body: Value = response.json().await.ok()?;
243        let token = body.get("token")?.as_str()?;
244        if let Some(challenge_url) = body.get("challenge_url").and_then(Value::as_str) {
245            if challenge_url.starts_with("https://") || challenge_url.starts_with("http://") {
246                return Some(challenge_url.to_string());
247            }
248        }
249
250        let mut platform_url = reqwest::Url::parse(self.config.api_url.trim_end_matches('/')).ok()?;
251        if let Some(host) = platform_url.host_str().map(str::to_string) {
252            if let Some(root_host) = host.strip_prefix("api.") {
253                platform_url.set_host(Some(root_host)).ok()?;
254            }
255        }
256        platform_url.set_path("");
257        platform_url.set_query(None);
258        platform_url.set_fragment(None);
259
260        Some(format!(
261            "{}/security-check?token={}",
262            platform_url.as_str().trim_end_matches('/'),
263            utf8_percent_encode(token, NON_ALPHANUMERIC),
264        ))
265    }
266
267    /// Evaluates the request against protection rules.
268    pub fn score_request(
269        &self,
270        ip: &str,
271        user_agent: Option<&str>,
272        headers: &HashMap<String, String>,
273        method: &str,
274        path: &str,
275    ) -> u32 {
276        let mut score: u32 = 0;
277
278        // UA evaluation
279        let ua = user_agent.unwrap_or("");
280        let ua_lower = ua.to_lowercase();
281        if ua.is_empty() {
282            score += 40;
283        } else {
284            let headless_keywords = ["puppeteer", "playwright", "phantomjs", "headlesschrome", "selenium"];
285            if headless_keywords.iter().any(|&k| ua_lower.contains(k)) {
286                score += 25;
287            }
288            let bot_keywords = ["googlebot", "bingbot", "yandex", "baiduspider", "curl", "wget", "httpclient", "python-urllib"];
289            if bot_keywords.iter().any(|&k| ua_lower.contains(k)) {
290                score += 35;
291            }
292        }
293
294        // Header check
295        if !headers.contains_key("accept") && !headers.contains_key("Accept") {
296            score += 15;
297        }
298
299        // Method Referrer check
300        if method.eq_ignore_ascii_case("POST") && !headers.contains_key("referer") && !headers.contains_key("Referer") {
301            score += 20;
302        }
303
304        // Rate Limit bucket check
305        if !self.consume_token(ip, path) {
306            score += 35;
307        }
308
309        let rules = self.rules.read().unwrap_or_else(|poisoned| poisoned.into_inner());
310        if let Some(custom_rules) = rules.as_ref().and_then(|value| value.get("rules")).and_then(Value::as_array) {
311            for rule in custom_rules {
312                let rule_type = rule.get("type").and_then(Value::as_str).unwrap_or("");
313                let pattern = rule.get("pattern").and_then(Value::as_str).unwrap_or("");
314                let condition = rule.get("condition").and_then(Value::as_str).unwrap_or("contains");
315                let candidate = match rule_type {
316                    "ip" => ip,
317                    "user_agent" => ua,
318                    "path" => path,
319                    _ => continue,
320                };
321                let matched = if condition == "equals" {
322                    candidate.eq_ignore_ascii_case(pattern)
323                } else {
324                    candidate.to_lowercase().contains(&pattern.to_lowercase())
325                };
326                if matched {
327                    score = score.saturating_add(rule.get("score").and_then(Value::as_u64).unwrap_or(0) as u32);
328                }
329            }
330        }
331
332        score.min(100)
333    }
334
335    fn consume_token(&self, ip: &str, path: &str) -> bool {
336        let now = unix_now() as f64;
337
338        let mut multiplier = 1.0;
339        let route_multipliers = [
340            ("/login", 0.4),
341            ("/auth", 0.4),
342            ("/api/", 0.7),
343            ("/assets/", 2.0),
344        ];
345
346        for (prefix, mult) in route_multipliers {
347            if path.starts_with(prefix) {
348                multiplier = mult;
349                break;
350            }
351        }
352
353        let burst = 24.0 * multiplier;
354        let rate_per_sec = 8.0 * multiplier;
355
356        let mut buckets = self.buckets.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
357        let bucket = buckets.entry(ip.to_string()).or_insert_with(|| TokenBucket {
358            tokens: burst,
359            last_ts: now,
360        });
361
362        let elapsed = now - bucket.last_ts;
363        bucket.last_ts = now;
364        bucket.tokens = (bucket.tokens + elapsed * rate_per_sec).min(burst);
365
366        if bucket.tokens >= 1.0 {
367            bucket.tokens -= 1.0;
368            true
369        } else {
370            false
371        }
372    }
373
374    /// Evaluates the full decision path.
375    pub async fn evaluate(
376        &self,
377        ip: &str,
378        user_agent: Option<&str>,
379        headers: &HashMap<String, String>,
380        method: &str,
381        path: &str,
382        domain: &str,
383    ) -> Decision {
384        // Sync rules if empty or expired
385        let now = unix_now();
386        let should_sync = now >= *self.next_sync_at.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
387
388        if should_sync && self.sync_in_progress.compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed).is_ok() {
389            let success = self.refresh_rules(domain).await.is_ok();
390            self.schedule_next_sync(success);
391            self.sync_in_progress.store(false, Ordering::Release);
392        }
393
394        let score = self.score_request(ip, user_agent, headers, method, path);
395
396        if score >= threshold("BLOCK") {
397            return Decision::Block;
398        }
399        if score >= threshold("DECOY") {
400            return Decision::Decoy;
401        }
402        if score >= threshold("CHALLENGE") {
403            return match self.challenge_url(domain, path).await {
404                Some(redirect_url) => Decision::Challenge { redirect_url },
405                None => Decision::Allow,
406            };
407        }
408        if score >= threshold("SLOW") {
409            return Decision::Slow;
410        }
411
412        Decision::Allow
413    }
414
415    fn schedule_next_sync(&self, success: bool) {
416        let mut failures = self.sync_failures.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
417        *failures = if success { 0 } else { failures.saturating_add(1).min(5) };
418        let base = if success {
419            self.config.sync_interval_seconds.max(10)
420        } else {
421            self.config.sync_interval_seconds.max(10).saturating_mul(1_u64 << *failures).min(300)
422        };
423        let jitter = SystemTime::now()
424            .duration_since(UNIX_EPOCH)
425            .map(|duration| 80 + (u64::from(duration.subsec_nanos()) % 41))
426            .unwrap_or(100);
427        let delay = (base.saturating_mul(jitter) / 100).max(8);
428        *self.next_sync_at.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) = unix_now().saturating_add(delay);
429    }
430
431    pub fn decoy_html() -> &'static str {
432        DECOY_HTML
433    }
434
435    pub fn block_html() -> &'static str {
436        BLOCK_HTML
437    }
438}
439
440fn unix_now() -> u64 {
441    SystemTime::now()
442        .duration_since(UNIX_EPOCH)
443        .map(|duration| duration.as_secs())
444        .unwrap_or(0)
445}
446
447fn threshold(tier: &str) -> u32 {
448    THRESHOLDS.iter()
449        .find_map(|(name, value)| (*name == tier).then_some(*value))
450        .unwrap_or(0)
451}