Skip to main content

sqlmap_rs/
client.rs

1//! Orchestrator for the `sqlmapapi.py` subprocess and its RESTful interface.
2//!
3//! Manages the full daemon lifecycle — boot, health check, task creation,
4//! scan execution, log retrieval, graceful stop/kill, and RAII cleanup.
5
6use crate::error::SqlmapError;
7use crate::types::{
8    BasicResponse, DataResponse, LogResponse, NewTaskResponse, SqlmapOptions, StatusResponse,
9};
10use reqwest::Client;
11use std::process::Stdio;
12use std::sync::atomic::{AtomicBool, Ordering};
13use std::sync::Arc;
14use std::time::Duration;
15use tokio::io::AsyncReadExt;
16use tokio::process::{Child, Command};
17use tokio::sync::Mutex;
18use tokio::time::sleep;
19use tracing::{debug, warn};
20
21const SPAWN_STDERR_SNIPPET_MAX: usize = 512;
22
23fn truncate_stderr_snippet(bytes: &[u8]) -> Option<String> {
24    if bytes.is_empty() {
25        return None;
26    }
27    let text = String::from_utf8_lossy(bytes);
28    let trimmed = text.trim();
29    if trimmed.is_empty() {
30        return None;
31    }
32    let snippet = if trimmed.len() > SPAWN_STDERR_SNIPPET_MAX {
33        format!("{}…", &trimmed[..SPAWN_STDERR_SNIPPET_MAX])
34    } else {
35        trimmed.to_string()
36    };
37    Some(snippet)
38}
39
40fn map_json_error(err: reqwest::Error) -> SqlmapError {
41    if err.is_decode() {
42        SqlmapError::MalformedResponse
43    } else {
44        SqlmapError::RequestError(err)
45    }
46}
47
48/// Manages the `sqlmapapi` lifecycle and provides access to its REST API.
49///
50/// When the engine is dropped, the locally spawned daemon subprocess
51/// is killed automatically via RAII (best-effort).
52pub struct SqlmapEngine {
53    api_url: String,
54    http: Client,
55    daemon_process: Option<Child>,
56    /// Configurable polling interval for `wait_for_completion`.
57    poll_interval: Duration,
58}
59
60impl SqlmapEngine {
61    /// Launches a local `sqlmapapi` daemon bound to `127.0.0.1`.
62    ///
63    /// # Arguments
64    ///
65    /// * `port` — TCP port for the daemon. Port `0` is not supported.
66    /// * `spawn_local` — If true, spawns a local `sqlmapapi` subprocess.
67    /// * `binary_path` — Override the `sqlmapapi` binary location.
68    ///
69    /// # Errors
70    ///
71    /// Returns [`SqlmapError::ProcessError`] if the daemon fails to spawn,
72    /// or [`SqlmapError::ApiError`] if it doesn't become responsive within 5 seconds.
73    pub async fn new(
74        port: u16,
75        spawn_local: bool,
76        binary_path: Option<&str>,
77    ) -> Result<Self, SqlmapError> {
78        Self::with_config(
79            port,
80            spawn_local,
81            binary_path,
82            Duration::from_secs(10),
83            Duration::from_millis(1000),
84        )
85        .await
86    }
87
88    /// Launches a daemon with custom HTTP timeout and polling interval.
89    ///
90    /// # Arguments
91    ///
92    /// * `request_timeout` — HTTP request timeout for API calls.
93    /// * `poll_interval` — Interval between status polls in `wait_for_completion`.
94    pub async fn with_config(
95        port: u16,
96        spawn_local: bool,
97        binary_path: Option<&str>,
98        request_timeout: Duration,
99        poll_interval: Duration,
100    ) -> Result<Self, SqlmapError> {
101        if poll_interval.is_zero() {
102            return Err(SqlmapError::ApiError(
103                "poll_interval must be greater than zero".into(),
104            ));
105        }
106
107        if request_timeout.is_zero() {
108            return Err(SqlmapError::ApiError(
109                "request_timeout must be greater than zero".into(),
110            ));
111        }
112
113        if port == 0 {
114            return Err(SqlmapError::ApiError("port 0 is not supported".into()));
115        }
116
117        let mut daemon_process = None;
118        let api_url = format!("http://127.0.0.1:{port}");
119
120        let http = Client::builder().timeout(request_timeout).build()?;
121
122        if spawn_local {
123            // Check if port is already in use before spawning.
124            if std::net::TcpStream::connect(format!("127.0.0.1:{port}")).is_ok() {
125                return Err(SqlmapError::PortConflict { port });
126            }
127
128            let binary = binary_path.unwrap_or("sqlmapapi");
129
130            let mut cmd = Command::new(binary);
131            cmd.arg("-s")
132                .arg("-H")
133                .arg("127.0.0.1")
134                .arg("-p")
135                .arg(port.to_string())
136                .kill_on_drop(true);
137
138            cmd.stdout(Stdio::null()).stderr(Stdio::piped());
139
140            let mut child = cmd.spawn().map_err(|e| {
141                if e.kind() == std::io::ErrorKind::NotFound {
142                    SqlmapError::BinaryNotFound(binary.to_string())
143                } else {
144                    SqlmapError::ProcessError(e)
145                }
146            })?;
147
148            let stderr_capture = Arc::new(Mutex::new(Vec::new()));
149            if let Some(mut stderr) = child.stderr.take() {
150                let capture = Arc::clone(&stderr_capture);
151                tokio::spawn(async move {
152                    let mut buf = vec![0u8; 4096];
153                    loop {
154                        match stderr.read(&mut buf).await {
155                            Ok(0) => break,
156                            Ok(n) => capture.lock().await.extend_from_slice(&buf[..n]),
157                            Err(_) => break,
158                        }
159                    }
160                });
161            }
162
163            daemon_process = Some(child);
164
165            // Wait for daemon to become responsive with a health probe.
166            let mut ready = false;
167            for attempt in 0..20 {
168                if Self::probe_task_new(&http, &api_url).await.is_ok() {
169                    ready = true;
170                    break;
171                }
172                debug!(attempt, "waiting for sqlmapapi daemon to become ready");
173                sleep(Duration::from_millis(250)).await;
174            }
175
176            if !ready {
177                let stderr_snippet = {
178                    let bytes = stderr_capture.lock().await;
179                    truncate_stderr_snippet(&bytes)
180                };
181                let mut msg =
182                    "sqlmapapi daemon failed to become responsive within 5 seconds".to_string();
183                if let Some(snippet) = stderr_snippet {
184                    msg.push_str(": ");
185                    msg.push_str(&snippet);
186                }
187                return Err(SqlmapError::ApiError(msg));
188            }
189        } else {
190            Self::probe_task_new(&http, &api_url).await?;
191        }
192
193        Ok(Self {
194            api_url,
195            http,
196            daemon_process,
197            poll_interval,
198        })
199    }
200
201    /// Probes `/task/new` and verifies the peer returns `success` + a non-empty `taskid`.
202    async fn probe_task_new(http: &Client, api_url: &str) -> Result<(), SqlmapError> {
203        let resp = http.get(format!("{api_url}/task/new")).send().await?;
204
205        if !resp.status().is_success() {
206            return Err(SqlmapError::ApiError(format!(
207                "health probe returned HTTP {}",
208                resp.status()
209            )));
210        }
211
212        let json = resp
213            .json::<NewTaskResponse>()
214            .await
215            .map_err(map_json_error)?;
216
217        if !json.success {
218            return Err(SqlmapError::ApiError(
219                json.message
220                    .unwrap_or_else(|| "health probe returned success=false".into()),
221            ));
222        }
223
224        let task_id = json.taskid.filter(|id| !id.is_empty()).ok_or_else(|| {
225            SqlmapError::ApiError("health probe: /task/new did not return a taskid".into())
226        })?;
227
228        let _ = http
229            .get(format!("{api_url}/task/{task_id}/delete"))
230            .send()
231            .await;
232
233        Ok(())
234    }
235
236    /// Creates and configures a new scanning task, returning an RAII wrapper.
237    ///
238    /// The task is automatically deleted from the daemon when dropped.
239    pub async fn create_task(
240        &self,
241        options: &SqlmapOptions,
242    ) -> Result<SqlmapTask<'_>, SqlmapError> {
243        let uri = format!("{}/task/new", self.api_url);
244        let resp = self.http.get(uri).send().await?;
245
246        if !resp.status().is_success() {
247            return Err(SqlmapError::ApiError(format!(
248                "task creation returned HTTP {}",
249                resp.status()
250            )));
251        }
252
253        let resp = resp
254            .json::<NewTaskResponse>()
255            .await
256            .map_err(map_json_error)?;
257
258        if !resp.success {
259            return Err(SqlmapError::ApiError(
260                resp.message
261                    .unwrap_or_else(|| "task creation returned success=false".into()),
262            ));
263        }
264
265        let task_id = match resp.taskid {
266            Some(id) if !id.is_empty() => id,
267            Some(id) => {
268                return Err(SqlmapError::InvalidTask(format!("empty taskid ({id:?})")));
269            }
270            None => {
271                return Err(SqlmapError::InvalidTask("missing taskid".into()));
272            }
273        };
274
275        let task = SqlmapTask {
276            engine: self,
277            task_id,
278            user_stopped: Arc::new(AtomicBool::new(false)),
279        };
280
281        // Set the configuration options on the new task.
282        let set_uri = format!("{}/option/{}/set", self.api_url, task.task_id);
283        let set_resp = self.http.post(&set_uri).json(options).send().await?;
284
285        if !set_resp.status().is_success() {
286            return Err(SqlmapError::ApiError(format!(
287                "option configuration returned HTTP {}",
288                set_resp.status()
289            )));
290        }
291
292        let set_resp = set_resp
293            .json::<BasicResponse>()
294            .await
295            .map_err(map_json_error)?;
296
297        if !set_resp.success {
298            return Err(SqlmapError::ApiError(
299                set_resp
300                    .message
301                    .unwrap_or_else(|| "option configuration failed".into()),
302            ));
303        }
304
305        Ok(task)
306    }
307
308    /// Check if sqlmapapi is available on this system.
309    ///
310    /// Tests that the `sqlmapapi` binary exists and is executable.
311    /// Does NOT fall back to `python3 -c "import sqlmap"` since that
312    /// doesn't guarantee the REST API server is available.
313    pub fn is_available() -> bool {
314        std::process::Command::new("sqlmapapi")
315            .arg("-h")
316            .stdout(Stdio::null())
317            .stderr(Stdio::null())
318            .status()
319            .map(|s| s.success())
320            .unwrap_or(false)
321    }
322
323    /// Check if sqlmapapi is available, trying the provided binary path first.
324    pub fn is_available_at(binary_path: &str) -> bool {
325        std::process::Command::new(binary_path)
326            .arg("-h")
327            .stdout(Stdio::null())
328            .stderr(Stdio::null())
329            .status()
330            .map(|s| s.success())
331            .unwrap_or(false)
332    }
333
334    /// Returns the base API URL for this engine.
335    pub fn api_url(&self) -> &str {
336        &self.api_url
337    }
338}
339
340impl Drop for SqlmapEngine {
341    fn drop(&mut self) {
342        if let Some(mut proc) = self.daemon_process.take() {
343            let _ = proc.start_kill();
344        }
345    }
346}
347
348// ── SqlmapTask ───────────────────────────────────────────────────
349
350/// An RAII-tracked scan execution task.
351///
352/// Ensures that the daemon reclaims task memory on drop by sending a
353/// delete request when a Tokio runtime is available (best-effort).
354pub struct SqlmapTask<'a> {
355    engine: &'a SqlmapEngine,
356    task_id: String,
357    /// Set after intentional [`stop`](Self::stop) or [`kill`](Self::kill).
358    user_stopped: Arc<AtomicBool>,
359}
360
361impl<'a> SqlmapTask<'a> {
362    /// Returns the unique task ID assigned by the daemon.
363    pub fn task_id(&self) -> &str {
364        &self.task_id
365    }
366
367    /// Starts the SQL injection scan on this task.
368    ///
369    /// The URL and options must have been configured via [`SqlmapEngine::create_task`].
370    pub async fn start(&self) -> Result<(), SqlmapError> {
371        let uri = format!("{}/scan/{}/start", self.engine.api_url, self.task_id);
372        let payload = serde_json::json!({});
373        let resp = self.engine.http.post(&uri).json(&payload).send().await?;
374
375        if resp.status().is_success() {
376            let body = resp.json::<BasicResponse>().await.map_err(map_json_error)?;
377            if !body.success {
378                return Err(SqlmapError::ApiError(
379                    body.message
380                        .unwrap_or_else(|| "scan start returned success=false".into()),
381                ));
382            }
383            Ok(())
384        } else {
385            Err(SqlmapError::ApiError(format!(
386                "scan start returned HTTP {}",
387                resp.status()
388            )))
389        }
390    }
391
392    /// Polls the task status until completion or timeout.
393    ///
394    /// Uses the engine's configured poll interval (default: 1 second).
395    pub async fn wait_for_completion(&self, timeout_secs: u64) -> Result<(), SqlmapError> {
396        let uri = format!("{}/scan/{}/status", self.engine.api_url, self.task_id);
397        let deadline = std::time::Instant::now() + Duration::from_secs(timeout_secs);
398
399        loop {
400            if std::time::Instant::now() >= deadline {
401                return Err(SqlmapError::Timeout(timeout_secs));
402            }
403
404            let resp = self.engine.http.get(&uri).send().await?;
405
406            if !resp.status().is_success() {
407                return Err(SqlmapError::ApiError(format!(
408                    "status check returned HTTP {}",
409                    resp.status()
410                )));
411            }
412
413            let status = resp
414                .json::<StatusResponse>()
415                .await
416                .map_err(map_json_error)?;
417
418            if !status.success {
419                return Err(SqlmapError::ApiError(
420                    status
421                        .message
422                        .unwrap_or_else(|| "status check returned success=false".into()),
423                ));
424            }
425
426            match status.status.as_deref() {
427                Some("running") => {
428                    debug!(task_id = %self.task_id, "scan running");
429                }
430                Some("terminated") => {
431                    if self.user_stopped.load(Ordering::Relaxed) {
432                        return Ok(());
433                    }
434                    match status.returncode {
435                        Some(0) => return Ok(()),
436                        Some(code) => {
437                            return Err(SqlmapError::ApiError(format!(
438                                "scan terminated with non-zero exit code {code}"
439                            )));
440                        }
441                        None => {
442                            return Err(SqlmapError::ApiError(
443                                "scan terminated without a process exit code".into(),
444                            ));
445                        }
446                    }
447                }
448                Some("not running") => {
449                    // sqlmapapi reports "not running" before start attaches a
450                    // process AND after some finished states. Keep polling;
451                    // only `terminated` is a definitive completion.
452                    debug!(task_id = %self.task_id, "scan not running yet");
453                }
454                Some(other) => {
455                    warn!(task_id = %self.task_id, status = %other, "unknown sqlmap status");
456                }
457                None => {}
458            }
459
460            sleep(self.engine.poll_interval).await;
461        }
462    }
463
464    /// Fetches the compiled data results from the engine.
465    pub async fn fetch_data(&self) -> Result<DataResponse, SqlmapError> {
466        let uri = format!("{}/scan/{}/data", self.engine.api_url, self.task_id);
467        let resp = self.engine.http.get(uri).send().await?;
468
469        if resp.status().is_success() {
470            let data = resp.json::<DataResponse>().await.map_err(map_json_error)?;
471            if !data.success {
472                return Err(SqlmapError::ApiError(
473                    data.message
474                        .unwrap_or_else(|| "data fetch returned success=false".into()),
475                ));
476            }
477            Ok(data)
478        } else {
479            Err(SqlmapError::ApiError(format!(
480                "data fetch returned HTTP {}",
481                resp.status()
482            )))
483        }
484    }
485
486    /// Fetches execution log entries for this task.
487    ///
488    /// Useful for monitoring what sqlmap is doing during a scan.
489    pub async fn fetch_log(&self) -> Result<LogResponse, SqlmapError> {
490        let uri = format!("{}/scan/{}/log", self.engine.api_url, self.task_id);
491        let resp = self.engine.http.get(uri).send().await?;
492
493        if resp.status().is_success() {
494            let log = resp.json::<LogResponse>().await.map_err(map_json_error)?;
495            if !log.success {
496                return Err(SqlmapError::ApiError(
497                    log.message
498                        .unwrap_or_else(|| "log fetch returned success=false".into()),
499                ));
500            }
501            Ok(log)
502        } else {
503            Err(SqlmapError::ApiError(format!(
504                "log fetch returned HTTP {}",
505                resp.status()
506            )))
507        }
508    }
509
510    /// Gracefully stops a running scan.
511    ///
512    /// The task can potentially be restarted after stopping.
513    pub async fn stop(&self) -> Result<(), SqlmapError> {
514        let uri = format!("{}/scan/{}/stop", self.engine.api_url, self.task_id);
515        let resp = self.engine.http.get(uri).send().await?;
516
517        if resp.status().is_success() {
518            let body = resp.json::<BasicResponse>().await.map_err(map_json_error)?;
519            if !body.success {
520                return Err(SqlmapError::ApiError(
521                    body.message
522                        .unwrap_or_else(|| "scan stop returned success=false".into()),
523                ));
524            }
525            self.user_stopped.store(true, Ordering::Relaxed);
526            Ok(())
527        } else {
528            Err(SqlmapError::ApiError(format!(
529                "scan stop returned HTTP {}",
530                resp.status()
531            )))
532        }
533    }
534
535    /// Forcefully kills a running scan.
536    ///
537    /// The task is terminated immediately. Data collected up to this point
538    /// may still be retrievable via [`fetch_data`](Self::fetch_data).
539    pub async fn kill(&self) -> Result<(), SqlmapError> {
540        let uri = format!("{}/scan/{}/kill", self.engine.api_url, self.task_id);
541        let resp = self.engine.http.get(uri).send().await?;
542
543        if resp.status().is_success() {
544            let body = resp.json::<BasicResponse>().await.map_err(map_json_error)?;
545            if !body.success {
546                return Err(SqlmapError::ApiError(
547                    body.message
548                        .unwrap_or_else(|| "scan kill returned success=false".into()),
549                ));
550            }
551            self.user_stopped.store(true, Ordering::Relaxed);
552            Ok(())
553        } else {
554            Err(SqlmapError::ApiError(format!(
555                "scan kill returned HTTP {}",
556                resp.status()
557            )))
558        }
559    }
560
561    /// Retrieves the current option values configured for this task.
562    pub async fn list_options(&self) -> Result<serde_json::Value, SqlmapError> {
563        let uri = format!("{}/option/{}/list", self.engine.api_url, self.task_id);
564        let resp = self.engine.http.get(uri).send().await?;
565
566        if resp.status().is_success() {
567            let value = resp
568                .json::<serde_json::Value>()
569                .await
570                .map_err(map_json_error)?;
571            if value
572                .get("success")
573                .and_then(|v| v.as_bool())
574                .is_some_and(|success| !success)
575            {
576                let message = value
577                    .get("message")
578                    .and_then(|v| v.as_str())
579                    .map(str::to_string)
580                    .unwrap_or_else(|| "option list returned success=false".into());
581                return Err(SqlmapError::ApiError(message));
582            }
583            Ok(value)
584        } else {
585            Err(SqlmapError::ApiError(format!(
586                "option list returned HTTP {}",
587                resp.status()
588            )))
589        }
590    }
591}
592
593impl<'a> Drop for SqlmapTask<'a> {
594    fn drop(&mut self) {
595        // Guarantee the server reclaims task memory when this struct goes out of scope.
596        // We use Handle::try_current() to avoid panicking if no Tokio runtime is active.
597        let uri = format!("{}/task/{}/delete", self.engine.api_url, self.task_id);
598        let client = self.engine.http.clone();
599
600        if let Ok(handle) = tokio::runtime::Handle::try_current() {
601            handle.spawn(async move {
602                let _ = client.get(&uri).send().await;
603            });
604        }
605        // If no runtime is available, we skip cleanup silently.
606        // The daemon will reclaim the task when it shuts down.
607    }
608}