From acc5b86aa48971157cb4aa55b17a7eed58bd9e6f Mon Sep 17 00:00:00 2001 From: Sharang Parnerkar Date: Thu, 12 Mar 2026 14:42:54 +0000 Subject: [PATCH] feat: AI-driven automated penetration testing (#12) --- Cargo.lock | 305 +++- Cargo.toml | 1 + compliance-agent/Cargo.toml | 1 + compliance-agent/src/api/handlers/mod.rs | 3 +- compliance-agent/src/api/handlers/pentest.rs | 718 ++++++++ compliance-agent/src/api/routes.rs | 38 + compliance-agent/src/database.rs | 45 + compliance-agent/src/llm/client.rs | 300 ++- compliance-agent/src/main.rs | 1 + compliance-agent/src/pentest/mod.rs | 5 + compliance-agent/src/pentest/orchestrator.rs | 761 ++++++++ compliance-agent/src/pentest/report.rs | 1601 +++++++++++++++++ compliance-core/src/models/dast.rs | 23 + compliance-core/src/models/mod.rs | 6 + compliance-core/src/models/pentest.rs | 294 +++ compliance-core/src/traits/mod.rs | 2 + compliance-core/src/traits/pentest_tool.rs | 63 + compliance-dashboard/assets/main.css | 464 +++++ compliance-dashboard/src/app.rs | 5 +- .../src/components/sidebar.rs | 6 + .../src/infrastructure/mod.rs | 1 + .../src/infrastructure/pentest.rs | 308 ++++ .../src/pages/dast_findings.rs | 104 +- compliance-dashboard/src/pages/mod.rs | 4 + .../src/pages/pentest_dashboard.rs | 398 ++++ .../src/pages/pentest_session.rs | 1141 ++++++++++++ compliance-dast/Cargo.toml | 4 + compliance-dast/src/lib.rs | 2 + compliance-dast/src/tools/api_fuzzer.rs | 146 ++ compliance-dast/src/tools/auth_bypass.rs | 130 ++ .../src/tools/console_log_detector.rs | 326 ++++ compliance-dast/src/tools/cookie_analyzer.rs | 401 +++++ compliance-dast/src/tools/cors_checker.rs | 410 +++++ compliance-dast/src/tools/csp_analyzer.rs | 447 +++++ compliance-dast/src/tools/dmarc_checker.rs | 401 +++++ compliance-dast/src/tools/dns_checker.rs | 389 ++++ compliance-dast/src/tools/mod.rs | 141 ++ compliance-dast/src/tools/openapi_parser.rs | 422 +++++ .../src/tools/rate_limit_tester.rs | 285 +++ compliance-dast/src/tools/recon.rs | 125 ++ compliance-dast/src/tools/security_headers.rs | 300 +++ compliance-dast/src/tools/sql_injection.rs | 138 ++ compliance-dast/src/tools/ssrf.rs | 134 ++ compliance-dast/src/tools/tls_analyzer.rs | 442 +++++ compliance-dast/src/tools/xss.rs | 134 ++ compliance-mcp/src/database.rs | 12 + compliance-mcp/src/server.rs | 52 +- compliance-mcp/src/tools/mod.rs | 1 + compliance-mcp/src/tools/pentest.rs | 261 +++ docs/.vitepress/config.mts | 1 + docs/features/dast.md | 15 +- docs/features/pentest.md | 110 ++ 52 files changed, 11729 insertions(+), 98 deletions(-) create mode 100644 compliance-agent/src/api/handlers/pentest.rs create mode 100644 compliance-agent/src/pentest/mod.rs create mode 100644 compliance-agent/src/pentest/orchestrator.rs create mode 100644 compliance-agent/src/pentest/report.rs create mode 100644 compliance-core/src/models/pentest.rs create mode 100644 compliance-core/src/traits/pentest_tool.rs create mode 100644 compliance-dashboard/src/infrastructure/pentest.rs create mode 100644 compliance-dashboard/src/pages/pentest_dashboard.rs create mode 100644 compliance-dashboard/src/pages/pentest_session.rs create mode 100644 compliance-dast/src/tools/api_fuzzer.rs create mode 100644 compliance-dast/src/tools/auth_bypass.rs create mode 100644 compliance-dast/src/tools/console_log_detector.rs create mode 100644 compliance-dast/src/tools/cookie_analyzer.rs create mode 100644 compliance-dast/src/tools/cors_checker.rs create mode 100644 compliance-dast/src/tools/csp_analyzer.rs create mode 100644 compliance-dast/src/tools/dmarc_checker.rs create mode 100644 compliance-dast/src/tools/dns_checker.rs create mode 100644 compliance-dast/src/tools/mod.rs create mode 100644 compliance-dast/src/tools/openapi_parser.rs create mode 100644 compliance-dast/src/tools/rate_limit_tester.rs create mode 100644 compliance-dast/src/tools/recon.rs create mode 100644 compliance-dast/src/tools/security_headers.rs create mode 100644 compliance-dast/src/tools/sql_injection.rs create mode 100644 compliance-dast/src/tools/ssrf.rs create mode 100644 compliance-dast/src/tools/tls_analyzer.rs create mode 100644 compliance-dast/src/tools/xss.rs create mode 100644 compliance-mcp/src/tools/pentest.rs create mode 100644 docs/features/pentest.md diff --git a/Cargo.lock b/Cargo.lock index d3e3e28..7a703ef 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,23 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures 0.2.17", +] + [[package]] name = "ahash" version = "0.8.12" @@ -45,6 +62,15 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[package]] +name = "arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" +dependencies = [ + "derive_arbitrary", +] + [[package]] name = "arc-swap" version = "1.8.2" @@ -391,6 +417,25 @@ dependencies = [ "serde", ] +[[package]] +name = "bzip2" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49ecfb22d906f800d4fe833b6282cf4dc1c298f5057ca0b5445e5c209735ca47" +dependencies = [ + "bzip2-sys", +] + +[[package]] +name = "bzip2-sys" +version = "0.1.13+1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "225bff33b2141874fe80d71e07d6eec4f85c5c216453dd96388240f96e1acc14" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "cc" version = "1.2.56" @@ -566,6 +611,16 @@ dependencies = [ "half", ] +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common", + "inout", +] + [[package]] name = "combine" version = "4.6.7" @@ -609,6 +664,7 @@ dependencies = [ "urlencoding", "uuid", "walkdir", + "zip", ] [[package]] @@ -680,12 +736,14 @@ dependencies = [ "chrono", "compliance-core", "mongodb", + "native-tls", "reqwest", "scraper", "serde", "serde_json", "thiserror 2.0.18", "tokio", + "tokio-native-tls", "tracing", "url", "uuid", @@ -834,6 +892,12 @@ dependencies = [ "unicode-xid", ] +[[package]] +name = "constant_time_eq" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c74b8349d32d297c9134b8c88677813a227df8f779daa29bfc29c183fe3dca6" + [[package]] name = "content_disposition" version = "0.4.0" @@ -939,6 +1003,21 @@ dependencies = [ "libc", ] +[[package]] +name = "crc" +version = "3.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5eb8a2a1cd12ab0d987a5d5e825195d372001a4094a0376319d5a0ad71c1ba0d" +dependencies = [ + "crc-catalog", +] + +[[package]] +name = "crc-catalog" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19d374276b40fb8bbdee95aef7c7fa6b5316ec764510eb64b8dd0e2ed0d7e7f5" + [[package]] name = "crc32fast" version = "1.5.0" @@ -1125,6 +1204,12 @@ version = "2.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d7a1e2f27636f116493b8b860f5546edb47c8d8f8ea73e1d2a20be88e28d1fea" +[[package]] +name = "deflate64" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "807800ff3288b621186fe0a8f3392c4652068257302709c24efd918c3dffcdc2" + [[package]] name = "deranged" version = "0.5.8" @@ -1157,6 +1242,17 @@ dependencies = [ "syn", ] +[[package]] +name = "derive_arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e567bd82dcff979e4b03460c307b3cdc9e96fde3d73bed1496d2bc75d9dd62a" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "derive_more" version = "2.1.1" @@ -1976,6 +2072,16 @@ version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + [[package]] name = "fnv" version = "1.0.7" @@ -1994,6 +2100,21 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" +[[package]] +name = "foreign-types" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" +dependencies = [ + "foreign-types-shared", +] + +[[package]] +name = "foreign-types-shared" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -2787,6 +2908,15 @@ dependencies = [ "cfb", ] +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "generic-array", +] + [[package]] name = "inventory" version = "0.3.22" @@ -2824,15 +2954,6 @@ dependencies = [ "serde", ] -[[package]] -name = "itertools" -version = "0.12.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" -dependencies = [ - "either", -] - [[package]] name = "itertools" version = "0.14.0" @@ -3076,6 +3197,27 @@ version = "0.11.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08ab2867e3eeeca90e844d1940eab391c9dc5228783db2ed999acbc0a9ed375a" +[[package]] +name = "lzma-rs" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "297e814c836ae64db86b36cf2a557ba54368d03f6afcd7d947c266692f71115e" +dependencies = [ + "byteorder", + "crc", +] + +[[package]] +name = "lzma-sys" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fda04ab3764e6cde78b9974eec4f779acaba7c4e84b36eca3cf77c581b85d27" +dependencies = [ + "cc", + "libc", + "pkg-config", +] + [[package]] name = "mac" version = "0.1.1" @@ -3272,6 +3414,16 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" version = "1.1.1" @@ -3399,6 +3551,23 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2195bf6aa996a481483b29d62a7663eed3fe39600c460e323f8ff41e90bdd89b" +[[package]] +name = "native-tls" +version = "0.2.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "465500e14ea162429d264d44189adc38b199b62b1c21eea9f69e4b73cb03bbf2" +dependencies = [ + "libc", + "log", + "openssl", + "openssl-probe 0.2.1", + "openssl-sys", + "schannel", + "security-framework", + "security-framework-sys", + "tempfile", +] + [[package]] name = "ndk" version = "0.9.0" @@ -3578,6 +3747,32 @@ version = "0.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "269bca4c2591a28585d6bf10d9ed0332b7d76900a1b02bec41bdc3a2cdcda107" +[[package]] +name = "openssl" +version = "0.10.75" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08838db121398ad17ab8531ce9de97b244589089e290a384c900cb9ff7434328" +dependencies = [ + "bitflags", + "cfg-if", + "foreign-types", + "libc", + "once_cell", + "openssl-macros", + "openssl-sys", +] + +[[package]] +name = "openssl-macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "openssl-probe" version = "0.1.6" @@ -3737,6 +3932,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2" dependencies = [ "digest", + "hmac", ] [[package]] @@ -3949,7 +4145,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.12.1", + "itertools", "proc-macro2", "quote", "syn", @@ -4806,6 +5002,12 @@ dependencies = [ "libc", ] +[[package]] +name = "simd-adler32" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e320a6c5ad31d271ad523dcf3ad13e2767ad8b1cb8f047f75a8aeaf8da139da2" + [[package]] name = "simple_asn1" version = "0.6.4" @@ -4899,7 +5101,7 @@ version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c1c97747dbf44bb1ca44a561ece23508e99cb592e862f22222dcf42f51d1e451" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn", @@ -5116,7 +5318,7 @@ dependencies = [ "fs4", "htmlescape", "hyperloglogplus", - "itertools 0.14.0", + "itertools", "levenshtein_automata", "log", "lru 0.12.5", @@ -5164,7 +5366,7 @@ checksum = "8b628488ae936c83e92b5c4056833054ca56f76c0e616aee8339e24ac89119cd" dependencies = [ "downcast-rs", "fastdivide", - "itertools 0.14.0", + "itertools", "serde", "tantivy-bitpacker", "tantivy-common", @@ -5214,7 +5416,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8292095d1a8a2c2b36380ec455f910ab52dde516af36321af332c93f20ab7d5" dependencies = [ "futures-util", - "itertools 0.14.0", + "itertools", "tantivy-bitpacker", "tantivy-common", "tantivy-fst", @@ -5428,6 +5630,16 @@ dependencies = [ "syn", ] +[[package]] +name = "tokio-native-tls" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbae76ab933c85776efabc971569dd6119c580d8f5d448769dec1764bf796ef2" +dependencies = [ + "native-tls", + "tokio", +] + [[package]] name = "tokio-rustls" version = "0.26.4" @@ -6805,6 +7017,15 @@ version = "0.8.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" +[[package]] +name = "xz2" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "388c44dc09d76f1536602ead6d325eb532f5c122f17782bd57fb47baeeb767e2" +dependencies = [ + "lzma-sys", +] + [[package]] name = "yoke" version = "0.8.1" @@ -6874,6 +7095,20 @@ name = "zeroize" version = "1.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85a5b4158499876c763cb03bc4e49185d3cccbabb15b33c627f7884f43db852e" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] [[package]] name = "zerotrie" @@ -6908,12 +7143,54 @@ dependencies = [ "syn", ] +[[package]] +name = "zip" +version = "2.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fabe6324e908f85a1c52063ce7aa26b68dcb7eb6dbc83a2d148403c9bc3eba50" +dependencies = [ + "aes", + "arbitrary", + "bzip2", + "constant_time_eq", + "crc32fast", + "crossbeam-utils", + "deflate64", + "displaydoc", + "flate2", + "getrandom 0.3.4", + "hmac", + "indexmap 2.13.0", + "lzma-rs", + "memchr", + "pbkdf2", + "sha1", + "thiserror 2.0.18", + "time", + "xz2", + "zeroize", + "zopfli", + "zstd", +] + [[package]] name = "zmij" version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" +[[package]] +name = "zopfli" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f05cd8797d63865425ff89b5c4a48804f35ba0ce8d125800027ad6017d2b5249" +dependencies = [ + "bumpalo", + "crc32fast", + "log", + "simd-adler32", +] + [[package]] name = "zstd" version = "0.13.3" diff --git a/Cargo.toml b/Cargo.toml index bc72548..34096b8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,3 +29,4 @@ hex = "0.4" uuid = { version = "1", features = ["v4", "serde"] } secrecy = { version = "0.10", features = ["serde"] } regex = "1" +zip = { version = "2", features = ["aes-crypto", "deflate"] } diff --git a/compliance-agent/Cargo.toml b/compliance-agent/Cargo.toml index a1e6eaf..bc2eac7 100644 --- a/compliance-agent/Cargo.toml +++ b/compliance-agent/Cargo.toml @@ -36,3 +36,4 @@ base64 = "0.22" urlencoding = "2" futures-util = "0.3" jsonwebtoken = "9" +zip = { workspace = true } diff --git a/compliance-agent/src/api/handlers/mod.rs b/compliance-agent/src/api/handlers/mod.rs index 8f353dd..63d64ef 100644 --- a/compliance-agent/src/api/handlers/mod.rs +++ b/compliance-agent/src/api/handlers/mod.rs @@ -1,6 +1,7 @@ pub mod chat; pub mod dast; pub mod graph; +pub mod pentest; use std::sync::Arc; @@ -1108,7 +1109,7 @@ pub async fn list_scan_runs( })) } -async fn collect_cursor_async( +pub(crate) async fn collect_cursor_async( mut cursor: mongodb::Cursor, ) -> Vec { use futures_util::StreamExt; diff --git a/compliance-agent/src/api/handlers/pentest.rs b/compliance-agent/src/api/handlers/pentest.rs new file mode 100644 index 0000000..57ecf4e --- /dev/null +++ b/compliance-agent/src/api/handlers/pentest.rs @@ -0,0 +1,718 @@ +use std::sync::Arc; + +use axum::extract::{Extension, Path, Query}; +use axum::http::StatusCode; +use axum::response::sse::{Event, Sse}; +use axum::response::IntoResponse; +use axum::Json; +use futures_util::stream; +use mongodb::bson::doc; +use serde::Deserialize; + +use compliance_core::models::dast::DastFinding; +use compliance_core::models::pentest::*; + +use crate::agent::ComplianceAgent; +use crate::pentest::PentestOrchestrator; + +use super::{collect_cursor_async, ApiResponse, PaginationParams}; + +type AgentExt = Extension>; + +#[derive(Deserialize)] +pub struct CreateSessionRequest { + pub target_id: String, + #[serde(default = "default_strategy")] + pub strategy: String, + pub message: Option, +} + +fn default_strategy() -> String { + "comprehensive".to_string() +} + +#[derive(Deserialize)] +pub struct SendMessageRequest { + pub message: String, +} + +/// POST /api/v1/pentest/sessions — Create a new pentest session and start the orchestrator +#[tracing::instrument(skip_all)] +pub async fn create_session( + Extension(agent): AgentExt, + Json(req): Json, +) -> Result>, (StatusCode, String)> { + let oid = mongodb::bson::oid::ObjectId::parse_str(&req.target_id).map_err(|_| { + ( + StatusCode::BAD_REQUEST, + "Invalid target_id format".to_string(), + ) + })?; + + // Look up the target + let target = agent + .db + .dast_targets() + .find_one(doc! { "_id": oid }) + .await + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Database error: {e}"), + ) + })? + .ok_or_else(|| (StatusCode::NOT_FOUND, "Target not found".to_string()))?; + + // Parse strategy + let strategy = match req.strategy.as_str() { + "quick" => PentestStrategy::Quick, + "targeted" => PentestStrategy::Targeted, + "aggressive" => PentestStrategy::Aggressive, + "stealth" => PentestStrategy::Stealth, + _ => PentestStrategy::Comprehensive, + }; + + // Create session + let mut session = PentestSession::new(req.target_id.clone(), strategy); + session.repo_id = target.repo_id.clone(); + + let insert_result = agent + .db + .pentest_sessions() + .insert_one(&session) + .await + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to create session: {e}"), + ) + })?; + + // Set the generated ID back on the session so the orchestrator has it + session.id = insert_result.inserted_id.as_object_id(); + + let initial_message = req.message.unwrap_or_else(|| { + format!( + "Begin a {} penetration test against {} ({}). \ + Identify vulnerabilities and provide evidence for each finding.", + session.strategy, target.name, target.base_url, + ) + }); + + // Spawn the orchestrator on a background task + let llm = agent.llm.clone(); + let db = agent.db.clone(); + let session_clone = session.clone(); + let target_clone = target.clone(); + tokio::spawn(async move { + let orchestrator = PentestOrchestrator::new(llm, db); + orchestrator + .run_session_guarded(&session_clone, &target_clone, &initial_message) + .await; + }); + + Ok(Json(ApiResponse { + data: session, + total: None, + page: None, + })) +} + +/// GET /api/v1/pentest/sessions — List pentest sessions +#[tracing::instrument(skip_all)] +pub async fn list_sessions( + Extension(agent): AgentExt, + Query(params): Query, +) -> Result>>, StatusCode> { + let db = &agent.db; + let skip = (params.page.saturating_sub(1)) * params.limit as u64; + let total = db + .pentest_sessions() + .count_documents(doc! {}) + .await + .unwrap_or(0); + + let sessions = match db + .pentest_sessions() + .find(doc! {}) + .sort(doc! { "started_at": -1 }) + .skip(skip) + .limit(params.limit) + .await + { + Ok(cursor) => collect_cursor_async(cursor).await, + Err(e) => { + tracing::warn!("Failed to fetch pentest sessions: {e}"); + Vec::new() + } + }; + + Ok(Json(ApiResponse { + data: sessions, + total: Some(total), + page: Some(params.page), + })) +} + +/// GET /api/v1/pentest/sessions/:id — Get a single pentest session +#[tracing::instrument(skip_all, fields(session_id = %id))] +pub async fn get_session( + Extension(agent): AgentExt, + Path(id): Path, +) -> Result>, StatusCode> { + let oid = + mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?; + + let session = agent + .db + .pentest_sessions() + .find_one(doc! { "_id": oid }) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + .ok_or(StatusCode::NOT_FOUND)?; + + Ok(Json(ApiResponse { + data: session, + total: None, + page: None, + })) +} + +/// POST /api/v1/pentest/sessions/:id/chat — Send a user message and trigger next orchestrator iteration +#[tracing::instrument(skip_all, fields(session_id = %id))] +pub async fn send_message( + Extension(agent): AgentExt, + Path(id): Path, + Json(req): Json, +) -> Result>, (StatusCode, String)> { + let oid = mongodb::bson::oid::ObjectId::parse_str(&id) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?; + + // Verify session exists and is running + let session = agent + .db + .pentest_sessions() + .find_one(doc! { "_id": oid }) + .await + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Database error: {e}"), + ) + })? + .ok_or_else(|| (StatusCode::NOT_FOUND, "Session not found".to_string()))?; + + if session.status != PentestStatus::Running && session.status != PentestStatus::Paused { + return Err(( + StatusCode::BAD_REQUEST, + format!("Session is {}, cannot send messages", session.status), + )); + } + + // Look up the target + let target_oid = + mongodb::bson::oid::ObjectId::parse_str(&session.target_id).map_err(|_| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + "Invalid target_id in session".to_string(), + ) + })?; + + let target = agent + .db + .dast_targets() + .find_one(doc! { "_id": target_oid }) + .await + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Database error: {e}"), + ) + })? + .ok_or_else(|| { + ( + StatusCode::NOT_FOUND, + "Target for session not found".to_string(), + ) + })?; + + // Store user message + let session_id = id.clone(); + let user_msg = PentestMessage::user(session_id.clone(), req.message.clone()); + let _ = agent.db.pentest_messages().insert_one(&user_msg).await; + + let response_msg = user_msg.clone(); + + // Spawn orchestrator to continue the session + let llm = agent.llm.clone(); + let db = agent.db.clone(); + let message = req.message.clone(); + tokio::spawn(async move { + let orchestrator = PentestOrchestrator::new(llm, db); + orchestrator + .run_session_guarded(&session, &target, &message) + .await; + }); + + Ok(Json(ApiResponse { + data: response_msg, + total: None, + page: None, + })) +} + +/// GET /api/v1/pentest/sessions/:id/stream — SSE endpoint for real-time events +/// +/// Returns recent messages as SSE events (polling approach). +/// True real-time streaming with broadcast channels will be added in a future iteration. +#[tracing::instrument(skip_all, fields(session_id = %id))] +pub async fn session_stream( + Extension(agent): AgentExt, + Path(id): Path, +) -> Result>>, StatusCode> +{ + let oid = + mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?; + + // Verify session exists + let _session = agent + .db + .pentest_sessions() + .find_one(doc! { "_id": oid }) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + .ok_or(StatusCode::NOT_FOUND)?; + + // Fetch recent messages for this session + let messages: Vec = match agent + .db + .pentest_messages() + .find(doc! { "session_id": &id }) + .sort(doc! { "created_at": 1 }) + .limit(100) + .await + { + Ok(cursor) => collect_cursor_async(cursor).await, + Err(_) => Vec::new(), + }; + + // Fetch recent attack chain nodes + let nodes: Vec = match agent + .db + .attack_chain_nodes() + .find(doc! { "session_id": &id }) + .sort(doc! { "started_at": 1 }) + .limit(100) + .await + { + Ok(cursor) => collect_cursor_async(cursor).await, + Err(_) => Vec::new(), + }; + + // Build SSE events from stored data + let mut events: Vec> = Vec::new(); + + for msg in &messages { + let event_data = serde_json::json!({ + "type": "message", + "role": msg.role, + "content": msg.content, + "created_at": msg.created_at.to_rfc3339(), + }); + if let Ok(data) = serde_json::to_string(&event_data) { + events.push(Ok(Event::default().event("message").data(data))); + } + } + + for node in &nodes { + let event_data = serde_json::json!({ + "type": "tool_execution", + "node_id": node.node_id, + "tool_name": node.tool_name, + "status": node.status, + "findings_produced": node.findings_produced, + }); + if let Ok(data) = serde_json::to_string(&event_data) { + events.push(Ok(Event::default().event("tool").data(data))); + } + } + + // Add session status event + let session = agent + .db + .pentest_sessions() + .find_one(doc! { "_id": oid }) + .await + .ok() + .flatten(); + + if let Some(s) = session { + let status_data = serde_json::json!({ + "type": "status", + "status": s.status, + "findings_count": s.findings_count, + "tool_invocations": s.tool_invocations, + }); + if let Ok(data) = serde_json::to_string(&status_data) { + events.push(Ok(Event::default().event("status").data(data))); + } + } + + Ok(Sse::new(stream::iter(events))) +} + +/// POST /api/v1/pentest/sessions/:id/stop — Stop a running pentest session +#[tracing::instrument(skip_all, fields(session_id = %id))] +pub async fn stop_session( + Extension(agent): AgentExt, + Path(id): Path, +) -> Result>, (StatusCode, String)> { + let oid = mongodb::bson::oid::ObjectId::parse_str(&id) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?; + + let session = agent + .db + .pentest_sessions() + .find_one(doc! { "_id": oid }) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Database error: {e}")))? + .ok_or_else(|| (StatusCode::NOT_FOUND, "Session not found".to_string()))?; + + if session.status != PentestStatus::Running { + return Err(( + StatusCode::BAD_REQUEST, + format!("Session is {}, not running", session.status), + )); + } + + agent + .db + .pentest_sessions() + .update_one( + doc! { "_id": oid }, + doc! { "$set": { + "status": "failed", + "completed_at": mongodb::bson::DateTime::now(), + "error_message": "Stopped by user", + }}, + ) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Database error: {e}")))?; + + let updated = agent + .db + .pentest_sessions() + .find_one(doc! { "_id": oid }) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Database error: {e}")))? + .ok_or_else(|| (StatusCode::NOT_FOUND, "Session not found after update".to_string()))?; + + Ok(Json(ApiResponse { + data: updated, + total: None, + page: None, + })) +} + +/// GET /api/v1/pentest/sessions/:id/attack-chain — Get attack chain nodes for a session +#[tracing::instrument(skip_all, fields(session_id = %id))] +pub async fn get_attack_chain( + Extension(agent): AgentExt, + Path(id): Path, +) -> Result>>, StatusCode> { + // Verify the session ID is valid + let _oid = + mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?; + + let nodes = match agent + .db + .attack_chain_nodes() + .find(doc! { "session_id": &id }) + .sort(doc! { "started_at": 1 }) + .await + { + Ok(cursor) => collect_cursor_async(cursor).await, + Err(e) => { + tracing::warn!("Failed to fetch attack chain nodes: {e}"); + Vec::new() + } + }; + + let total = nodes.len() as u64; + Ok(Json(ApiResponse { + data: nodes, + total: Some(total), + page: None, + })) +} + +/// GET /api/v1/pentest/sessions/:id/messages — Get messages for a session +#[tracing::instrument(skip_all, fields(session_id = %id))] +pub async fn get_messages( + Extension(agent): AgentExt, + Path(id): Path, + Query(params): Query, +) -> Result>>, StatusCode> { + let _oid = + mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?; + + let skip = (params.page.saturating_sub(1)) * params.limit as u64; + let total = agent + .db + .pentest_messages() + .count_documents(doc! { "session_id": &id }) + .await + .unwrap_or(0); + + let messages = match agent + .db + .pentest_messages() + .find(doc! { "session_id": &id }) + .sort(doc! { "created_at": 1 }) + .skip(skip) + .limit(params.limit) + .await + { + Ok(cursor) => collect_cursor_async(cursor).await, + Err(e) => { + tracing::warn!("Failed to fetch pentest messages: {e}"); + Vec::new() + } + }; + + Ok(Json(ApiResponse { + data: messages, + total: Some(total), + page: Some(params.page), + })) +} + +/// GET /api/v1/pentest/stats — Aggregated pentest statistics +#[tracing::instrument(skip_all)] +pub async fn pentest_stats( + Extension(agent): AgentExt, +) -> Result>, StatusCode> { + let db = &agent.db; + + let running_sessions = db + .pentest_sessions() + .count_documents(doc! { "status": "running" }) + .await + .unwrap_or(0) as u32; + + // Count DAST findings from pentest sessions + let total_vulnerabilities = db + .dast_findings() + .count_documents(doc! { "session_id": { "$exists": true, "$ne": null } }) + .await + .unwrap_or(0) as u32; + + // Aggregate tool invocations from all sessions + let sessions: Vec = match db.pentest_sessions().find(doc! {}).await { + Ok(cursor) => collect_cursor_async(cursor).await, + Err(_) => Vec::new(), + }; + + let total_tool_invocations: u32 = sessions.iter().map(|s| s.tool_invocations).sum(); + let total_successes: u32 = sessions.iter().map(|s| s.tool_successes).sum(); + let tool_success_rate = if total_tool_invocations == 0 { + 100.0 + } else { + (total_successes as f64 / total_tool_invocations as f64) * 100.0 + }; + + // Severity distribution from pentest-related DAST findings + let critical = db + .dast_findings() + .count_documents(doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "critical" }) + .await + .unwrap_or(0) as u32; + let high = db + .dast_findings() + .count_documents(doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "high" }) + .await + .unwrap_or(0) as u32; + let medium = db + .dast_findings() + .count_documents(doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "medium" }) + .await + .unwrap_or(0) as u32; + let low = db + .dast_findings() + .count_documents(doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "low" }) + .await + .unwrap_or(0) as u32; + let info = db + .dast_findings() + .count_documents(doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "info" }) + .await + .unwrap_or(0) as u32; + + Ok(Json(ApiResponse { + data: PentestStats { + running_sessions, + total_vulnerabilities, + total_tool_invocations, + tool_success_rate, + severity_distribution: SeverityDistribution { + critical, + high, + medium, + low, + info, + }, + }, + total: None, + page: None, + })) +} + +/// GET /api/v1/pentest/sessions/:id/findings — Get DAST findings for a pentest session +#[tracing::instrument(skip_all, fields(session_id = %id))] +pub async fn get_session_findings( + Extension(agent): AgentExt, + Path(id): Path, + Query(params): Query, +) -> Result>>, StatusCode> { + let _oid = + mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?; + + let skip = (params.page.saturating_sub(1)) * params.limit as u64; + let total = agent + .db + .dast_findings() + .count_documents(doc! { "session_id": &id }) + .await + .unwrap_or(0); + + let findings = match agent + .db + .dast_findings() + .find(doc! { "session_id": &id }) + .sort(doc! { "created_at": -1 }) + .skip(skip) + .limit(params.limit) + .await + { + Ok(cursor) => collect_cursor_async(cursor).await, + Err(e) => { + tracing::warn!("Failed to fetch pentest session findings: {e}"); + Vec::new() + } + }; + + Ok(Json(ApiResponse { + data: findings, + total: Some(total), + page: Some(params.page), + })) +} + +#[derive(Deserialize)] +pub struct ExportBody { + pub password: String, + /// Requester display name (from auth) + #[serde(default)] + pub requester_name: String, + /// Requester email (from auth) + #[serde(default)] + pub requester_email: String, +} + +/// POST /api/v1/pentest/sessions/:id/export — Export an encrypted pentest report archive +#[tracing::instrument(skip_all, fields(session_id = %id))] +pub async fn export_session_report( + Extension(agent): AgentExt, + Path(id): Path, + Json(body): Json, +) -> Result { + let oid = mongodb::bson::oid::ObjectId::parse_str(&id) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?; + + if body.password.len() < 8 { + return Err(( + StatusCode::BAD_REQUEST, + "Password must be at least 8 characters".to_string(), + )); + } + + // Fetch session + let session = agent + .db + .pentest_sessions() + .find_one(doc! { "_id": oid }) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Database error: {e}")))? + .ok_or_else(|| (StatusCode::NOT_FOUND, "Session not found".to_string()))?; + + // Resolve target name + let target = if let Ok(tid) = mongodb::bson::oid::ObjectId::parse_str(&session.target_id) { + agent + .db + .dast_targets() + .find_one(doc! { "_id": tid }) + .await + .ok() + .flatten() + } else { + None + }; + let target_name = target + .as_ref() + .map(|t| t.name.clone()) + .unwrap_or_else(|| "Unknown Target".to_string()); + let target_url = target + .as_ref() + .map(|t| t.base_url.clone()) + .unwrap_or_default(); + + // Fetch attack chain nodes + let nodes: Vec = match agent + .db + .attack_chain_nodes() + .find(doc! { "session_id": &id }) + .sort(doc! { "started_at": 1 }) + .await + { + Ok(cursor) => collect_cursor_async(cursor).await, + Err(_) => Vec::new(), + }; + + // Fetch DAST findings for this session + let findings: Vec = match agent + .db + .dast_findings() + .find(doc! { "session_id": &id }) + .sort(doc! { "severity": -1, "created_at": -1 }) + .await + { + Ok(cursor) => collect_cursor_async(cursor).await, + Err(_) => Vec::new(), + }; + + let ctx = crate::pentest::report::ReportContext { + session, + target_name, + target_url, + findings, + attack_chain: nodes, + requester_name: if body.requester_name.is_empty() { + "Unknown".to_string() + } else { + body.requester_name + }, + requester_email: body.requester_email, + }; + + let report = crate::pentest::generate_encrypted_report(&ctx, &body.password) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?; + + let response = serde_json::json!({ + "archive_base64": base64::Engine::encode(&base64::engine::general_purpose::STANDARD, &report.archive), + "sha256": report.sha256, + "filename": format!("pentest-report-{id}.zip"), + }); + + Ok(Json(response).into_response()) +} diff --git a/compliance-agent/src/api/routes.rs b/compliance-agent/src/api/routes.rs index 502984b..f878d7e 100644 --- a/compliance-agent/src/api/routes.rs +++ b/compliance-agent/src/api/routes.rs @@ -99,6 +99,44 @@ pub fn build_router() -> Router { "/api/v1/chat/{repo_id}/status", get(handlers::chat::embedding_status), ) + // Pentest API endpoints + .route( + "/api/v1/pentest/sessions", + get(handlers::pentest::list_sessions).post(handlers::pentest::create_session), + ) + .route( + "/api/v1/pentest/sessions/{id}", + get(handlers::pentest::get_session), + ) + .route( + "/api/v1/pentest/sessions/{id}/chat", + post(handlers::pentest::send_message), + ) + .route( + "/api/v1/pentest/sessions/{id}/stop", + post(handlers::pentest::stop_session), + ) + .route( + "/api/v1/pentest/sessions/{id}/stream", + get(handlers::pentest::session_stream), + ) + .route( + "/api/v1/pentest/sessions/{id}/attack-chain", + get(handlers::pentest::get_attack_chain), + ) + .route( + "/api/v1/pentest/sessions/{id}/messages", + get(handlers::pentest::get_messages), + ) + .route( + "/api/v1/pentest/sessions/{id}/findings", + get(handlers::pentest::get_session_findings), + ) + .route( + "/api/v1/pentest/sessions/{id}/export", + post(handlers::pentest::export_session_report), + ) + .route("/api/v1/pentest/stats", get(handlers::pentest::pentest_stats)) // Webhook endpoints (proxied through dashboard) .route( "/webhook/github/{repo_id}", diff --git a/compliance-agent/src/database.rs b/compliance-agent/src/database.rs index c2b0740..6b0c0d9 100644 --- a/compliance-agent/src/database.rs +++ b/compliance-agent/src/database.rs @@ -166,6 +166,38 @@ impl Database { ) .await?; + // pentest_sessions: compound (target_id, started_at DESC) + self.pentest_sessions() + .create_index( + IndexModel::builder() + .keys(doc! { "target_id": 1, "started_at": -1 }) + .build(), + ) + .await?; + + // pentest_sessions: status index + self.pentest_sessions() + .create_index(IndexModel::builder().keys(doc! { "status": 1 }).build()) + .await?; + + // attack_chain_nodes: compound (session_id, node_id) + self.attack_chain_nodes() + .create_index( + IndexModel::builder() + .keys(doc! { "session_id": 1, "node_id": 1 }) + .build(), + ) + .await?; + + // pentest_messages: compound (session_id, created_at) + self.pentest_messages() + .create_index( + IndexModel::builder() + .keys(doc! { "session_id": 1, "created_at": 1 }) + .build(), + ) + .await?; + tracing::info!("Database indexes ensured"); Ok(()) } @@ -235,6 +267,19 @@ impl Database { self.inner.collection("embedding_builds") } + // Pentest collections + pub fn pentest_sessions(&self) -> Collection { + self.inner.collection("pentest_sessions") + } + + pub fn attack_chain_nodes(&self) -> Collection { + self.inner.collection("attack_chain_nodes") + } + + pub fn pentest_messages(&self) -> Collection { + self.inner.collection("pentest_messages") + } + #[allow(dead_code)] pub fn raw_collection(&self, name: &str) -> Collection { self.inner.collection(name) diff --git a/compliance-agent/src/llm/client.rs b/compliance-agent/src/llm/client.rs index c7a571d..b6bc657 100644 --- a/compliance-agent/src/llm/client.rs +++ b/compliance-agent/src/llm/client.rs @@ -12,10 +12,16 @@ pub struct LlmClient { http: reqwest::Client, } -#[derive(Serialize)] -struct ChatMessage { - role: String, - content: String, +// ── Request types ────────────────────────────────────────────── + +#[derive(Serialize, Clone, Debug)] +pub struct ChatMessage { + pub role: String, + pub content: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_call_id: Option, } #[derive(Serialize)] @@ -26,8 +32,25 @@ struct ChatCompletionRequest { temperature: Option, #[serde(skip_serializing_if = "Option::is_none")] max_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + tools: Option>, } +#[derive(Serialize)] +struct ToolDefinitionPayload { + r#type: String, + function: ToolFunctionPayload, +} + +#[derive(Serialize)] +struct ToolFunctionPayload { + name: String, + description: String, + parameters: serde_json::Value, +} + +// ── Response types ───────────────────────────────────────────── + #[derive(Deserialize)] struct ChatCompletionResponse { choices: Vec, @@ -40,29 +63,85 @@ struct ChatChoice { #[derive(Deserialize)] struct ChatResponseMessage { - content: String, + #[serde(default)] + content: Option, + #[serde(default)] + tool_calls: Option>, } -/// Request body for the embeddings API +#[derive(Deserialize)] +struct ToolCallResponse { + id: String, + function: ToolCallFunction, +} + +#[derive(Deserialize)] +struct ToolCallFunction { + name: String, + arguments: String, +} + +// ── Public types for tool calling ────────────────────────────── + +/// Definition of a tool that the LLM can invoke +#[derive(Debug, Clone, Serialize)] +pub struct ToolDefinition { + pub name: String, + pub description: String, + pub parameters: serde_json::Value, +} + +/// A tool call request from the LLM +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct LlmToolCall { + pub id: String, + pub name: String, + pub arguments: serde_json::Value, +} + +/// A tool call in the request message format (for sending back tool_calls in assistant messages) +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolCallRequest { + pub id: String, + pub r#type: String, + pub function: ToolCallRequestFunction, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolCallRequestFunction { + pub name: String, + pub arguments: String, +} + +/// Response from the LLM — either content or tool calls +#[derive(Debug, Clone)] +pub enum LlmResponse { + Content(String), + /// Tool calls with optional reasoning text from the LLM + ToolCalls { calls: Vec, reasoning: String }, +} + +// ── Embedding types ──────────────────────────────────────────── + #[derive(Serialize)] struct EmbeddingRequest { model: String, input: Vec, } -/// Response from the embeddings API #[derive(Deserialize)] struct EmbeddingResponse { data: Vec, } -/// A single embedding result #[derive(Deserialize)] struct EmbeddingData { embedding: Vec, index: usize, } +// ── Implementation ───────────────────────────────────────────── + impl LlmClient { pub fn new( base_url: String, @@ -83,98 +162,142 @@ impl LlmClient { &self.embed_model } + fn chat_url(&self) -> String { + format!( + "{}/v1/chat/completions", + self.base_url.trim_end_matches('/') + ) + } + + fn auth_header(&self) -> Option { + let key = self.api_key.expose_secret(); + if key.is_empty() { + None + } else { + Some(format!("Bearer {key}")) + } + } + + /// Simple chat: system + user prompt → text response pub async fn chat( &self, system_prompt: &str, user_prompt: &str, temperature: Option, ) -> Result { - let url = format!( - "{}/v1/chat/completions", - self.base_url.trim_end_matches('/') - ); + let messages = vec![ + ChatMessage { + role: "system".to_string(), + content: Some(system_prompt.to_string()), + tool_calls: None, + tool_call_id: None, + }, + ChatMessage { + role: "user".to_string(), + content: Some(user_prompt.to_string()), + tool_calls: None, + tool_call_id: None, + }, + ]; let request_body = ChatCompletionRequest { model: self.model.clone(), - messages: vec![ - ChatMessage { - role: "system".to_string(), - content: system_prompt.to_string(), - }, - ChatMessage { - role: "user".to_string(), - content: user_prompt.to_string(), - }, - ], + messages, temperature, max_tokens: Some(4096), + tools: None, }; - let mut req = self - .http - .post(&url) - .header("content-type", "application/json") - .json(&request_body); - - let key = self.api_key.expose_secret(); - if !key.is_empty() { - req = req.header("Authorization", format!("Bearer {key}")); - } - - let resp = req - .send() - .await - .map_err(|e| AgentError::Other(format!("LiteLLM request failed: {e}")))?; - - if !resp.status().is_success() { - let status = resp.status(); - let body = resp.text().await.unwrap_or_default(); - return Err(AgentError::Other(format!( - "LiteLLM returned {status}: {body}" - ))); - } - - let body: ChatCompletionResponse = resp - .json() - .await - .map_err(|e| AgentError::Other(format!("Failed to parse LiteLLM response: {e}")))?; - - body.choices - .first() - .map(|c| c.message.content.clone()) - .ok_or_else(|| AgentError::Other("Empty response from LiteLLM".to_string())) + self.send_chat_request(&request_body).await.map(|resp| { + match resp { + LlmResponse::Content(c) => c, + LlmResponse::ToolCalls { .. } => String::new(), // shouldn't happen without tools + } + }) } + /// Chat with a list of (role, content) messages → text response #[allow(dead_code)] pub async fn chat_with_messages( &self, messages: Vec<(String, String)>, temperature: Option, ) -> Result { - let url = format!( - "{}/v1/chat/completions", - self.base_url.trim_end_matches('/') - ); + let messages = messages + .into_iter() + .map(|(role, content)| ChatMessage { + role, + content: Some(content), + tool_calls: None, + tool_call_id: None, + }) + .collect(); let request_body = ChatCompletionRequest { model: self.model.clone(), - messages: messages - .into_iter() - .map(|(role, content)| ChatMessage { role, content }) - .collect(), + messages, temperature, max_tokens: Some(4096), + tools: None, }; + self.send_chat_request(&request_body).await.map(|resp| { + match resp { + LlmResponse::Content(c) => c, + LlmResponse::ToolCalls { .. } => String::new(), + } + }) + } + + /// Chat with tool definitions — returns either content or tool calls. + /// Use this for the AI pentest orchestrator loop. + pub async fn chat_with_tools( + &self, + messages: Vec, + tools: &[ToolDefinition], + temperature: Option, + max_tokens: Option, + ) -> Result { + let tool_payloads: Vec = tools + .iter() + .map(|t| ToolDefinitionPayload { + r#type: "function".to_string(), + function: ToolFunctionPayload { + name: t.name.clone(), + description: t.description.clone(), + parameters: t.parameters.clone(), + }, + }) + .collect(); + + let request_body = ChatCompletionRequest { + model: self.model.clone(), + messages, + temperature, + max_tokens: Some(max_tokens.unwrap_or(8192)), + tools: if tool_payloads.is_empty() { + None + } else { + Some(tool_payloads) + }, + }; + + self.send_chat_request(&request_body).await + } + + /// Internal method to send a chat completion request and parse the response + async fn send_chat_request( + &self, + request_body: &ChatCompletionRequest, + ) -> Result { let mut req = self .http - .post(&url) + .post(&self.chat_url()) .header("content-type", "application/json") - .json(&request_body); + .json(request_body); - let key = self.api_key.expose_secret(); - if !key.is_empty() { - req = req.header("Authorization", format!("Bearer {key}")); + if let Some(auth) = self.auth_header() { + req = req.header("Authorization", auth); } let resp = req @@ -195,10 +318,39 @@ impl LlmClient { .await .map_err(|e| AgentError::Other(format!("Failed to parse LiteLLM response: {e}")))?; - body.choices + let choice = body + .choices .first() - .map(|c| c.message.content.clone()) - .ok_or_else(|| AgentError::Other("Empty response from LiteLLM".to_string())) + .ok_or_else(|| AgentError::Other("Empty response from LiteLLM".to_string()))?; + + // Check for tool calls first + if let Some(tool_calls) = &choice.message.tool_calls { + if !tool_calls.is_empty() { + let calls: Vec = tool_calls + .iter() + .map(|tc| { + let arguments = serde_json::from_str(&tc.function.arguments) + .unwrap_or(serde_json::Value::Object(serde_json::Map::new())); + LlmToolCall { + id: tc.id.clone(), + name: tc.function.name.clone(), + arguments, + } + }) + .collect(); + // Capture any reasoning text the LLM included alongside tool calls + let reasoning = choice.message.content.clone().unwrap_or_default(); + return Ok(LlmResponse::ToolCalls { calls, reasoning }); + } + } + + // Otherwise return content + let content = choice + .message + .content + .clone() + .unwrap_or_default(); + Ok(LlmResponse::Content(content)) } /// Generate embeddings for a batch of texts @@ -216,9 +368,8 @@ impl LlmClient { .header("content-type", "application/json") .json(&request_body); - let key = self.api_key.expose_secret(); - if !key.is_empty() { - req = req.header("Authorization", format!("Bearer {key}")); + if let Some(auth) = self.auth_header() { + req = req.header("Authorization", auth); } let resp = req @@ -239,7 +390,6 @@ impl LlmClient { .await .map_err(|e| AgentError::Other(format!("Failed to parse embedding response: {e}")))?; - // Sort by index to maintain input order let mut data = body.data; data.sort_by_key(|d| d.index); diff --git a/compliance-agent/src/main.rs b/compliance-agent/src/main.rs index 97ec23e..9e346e9 100644 --- a/compliance-agent/src/main.rs +++ b/compliance-agent/src/main.rs @@ -4,6 +4,7 @@ mod config; mod database; mod error; mod llm; +mod pentest; mod pipeline; mod rag; mod scheduler; diff --git a/compliance-agent/src/pentest/mod.rs b/compliance-agent/src/pentest/mod.rs new file mode 100644 index 0000000..934315a --- /dev/null +++ b/compliance-agent/src/pentest/mod.rs @@ -0,0 +1,5 @@ +pub mod orchestrator; +pub mod report; + +pub use orchestrator::PentestOrchestrator; +pub use report::generate_encrypted_report; diff --git a/compliance-agent/src/pentest/orchestrator.rs b/compliance-agent/src/pentest/orchestrator.rs new file mode 100644 index 0000000..184db47 --- /dev/null +++ b/compliance-agent/src/pentest/orchestrator.rs @@ -0,0 +1,761 @@ +use std::sync::Arc; +use std::time::Duration; + +use futures_util::StreamExt; +use mongodb::bson::doc; +use tokio::sync::broadcast; + +use compliance_core::models::dast::DastTarget; +use compliance_core::models::finding::{Finding, FindingStatus, Severity}; +use compliance_core::models::pentest::*; +use compliance_core::models::sbom::SbomEntry; +use compliance_core::traits::pentest_tool::PentestToolContext; +use compliance_dast::ToolRegistry; + +use crate::database::Database; +use crate::llm::client::{ + ChatMessage, LlmResponse, ToolCallRequest, ToolCallRequestFunction, ToolDefinition, +}; +use crate::llm::LlmClient; + +/// Maximum duration for a single pentest session before timeout +const SESSION_TIMEOUT: Duration = Duration::from_secs(30 * 60); // 30 minutes + +pub struct PentestOrchestrator { + tool_registry: ToolRegistry, + llm: Arc, + db: Database, + event_tx: broadcast::Sender, +} + +impl PentestOrchestrator { + pub fn new(llm: Arc, db: Database) -> Self { + let (event_tx, _) = broadcast::channel(256); + Self { + tool_registry: ToolRegistry::new(), + llm, + db, + event_tx, + } + } + + pub fn subscribe(&self) -> broadcast::Receiver { + self.event_tx.subscribe() + } + + pub fn event_sender(&self) -> broadcast::Sender { + self.event_tx.clone() + } + + /// Run a pentest session with timeout and automatic failure marking on errors. + pub async fn run_session_guarded( + &self, + session: &PentestSession, + target: &DastTarget, + initial_message: &str, + ) { + let session_id = session.id; + + match tokio::time::timeout( + SESSION_TIMEOUT, + self.run_session(session, target, initial_message), + ) + .await + { + Ok(Ok(())) => { + tracing::info!(?session_id, "Pentest session completed successfully"); + } + Ok(Err(e)) => { + tracing::error!(?session_id, error = %e, "Pentest session failed"); + self.mark_session_failed(session_id, &format!("Error: {e}")) + .await; + let _ = self.event_tx.send(PentestEvent::Error { + message: format!("Session failed: {e}"), + }); + } + Err(_) => { + tracing::warn!(?session_id, "Pentest session timed out after 30 minutes"); + self.mark_session_failed(session_id, "Session timed out after 30 minutes") + .await; + let _ = self.event_tx.send(PentestEvent::Error { + message: "Session timed out after 30 minutes".to_string(), + }); + } + } + } + + async fn mark_session_failed( + &self, + session_id: Option, + reason: &str, + ) { + if let Some(sid) = session_id { + let _ = self + .db + .pentest_sessions() + .update_one( + doc! { "_id": sid }, + doc! { "$set": { + "status": "failed", + "completed_at": mongodb::bson::DateTime::now(), + "error_message": reason, + }}, + ) + .await; + } + } + + async fn run_session( + &self, + session: &PentestSession, + target: &DastTarget, + initial_message: &str, + ) -> Result<(), crate::error::AgentError> { + let session_id = session + .id + .map(|oid| oid.to_hex()) + .unwrap_or_default(); + + // Gather code-awareness context from linked repo + let (sast_findings, sbom_entries, code_context) = + self.gather_repo_context(target).await; + + // Build system prompt with code context + let system_prompt = self + .build_system_prompt(session, target, &sast_findings, &sbom_entries, &code_context) + .await; + + // Build tool definitions for LLM + let tool_defs: Vec = self + .tool_registry + .all_definitions() + .into_iter() + .map(|td| ToolDefinition { + name: td.name, + description: td.description, + parameters: td.input_schema, + }) + .collect(); + + // Initialize messages + let mut messages = vec![ + ChatMessage { + role: "system".to_string(), + content: Some(system_prompt), + tool_calls: None, + tool_call_id: None, + }, + ChatMessage { + role: "user".to_string(), + content: Some(initial_message.to_string()), + tool_calls: None, + tool_call_id: None, + }, + ]; + + // Store user message + let user_msg = PentestMessage::user(session_id.clone(), initial_message.to_string()); + let _ = self.db.pentest_messages().insert_one(&user_msg).await; + + // Build tool context with real data + let tool_context = PentestToolContext { + target: target.clone(), + session_id: session_id.clone(), + sast_findings, + sbom_entries, + code_context, + rate_limit: target.rate_limit, + allow_destructive: target.allow_destructive, + }; + + let max_iterations = 50; + let mut total_findings = 0u32; + let mut total_tool_calls = 0u32; + let mut total_successes = 0u32; + let mut prev_node_ids: Vec = Vec::new(); + + for _iteration in 0..max_iterations { + let response = self + .llm + .chat_with_tools(messages.clone(), &tool_defs, Some(0.2), Some(8192)) + .await?; + + match response { + LlmResponse::Content(content) => { + let msg = + PentestMessage::assistant(session_id.clone(), content.clone()); + let _ = self.db.pentest_messages().insert_one(&msg).await; + let _ = self.event_tx.send(PentestEvent::Message { + content: content.clone(), + }); + + messages.push(ChatMessage { + role: "assistant".to_string(), + content: Some(content.clone()), + tool_calls: None, + tool_call_id: None, + }); + + let done_indicators = [ + "pentest complete", + "testing complete", + "scan complete", + "analysis complete", + "finished", + "that concludes", + ]; + let content_lower = content.to_lowercase(); + if done_indicators + .iter() + .any(|ind| content_lower.contains(ind)) + { + break; + } + break; + } + LlmResponse::ToolCalls { calls: tool_calls, reasoning } => { + let tc_requests: Vec = tool_calls + .iter() + .map(|tc| ToolCallRequest { + id: tc.id.clone(), + r#type: "function".to_string(), + function: ToolCallRequestFunction { + name: tc.name.clone(), + arguments: serde_json::to_string(&tc.arguments) + .unwrap_or_default(), + }, + }) + .collect(); + + messages.push(ChatMessage { + role: "assistant".to_string(), + content: if reasoning.is_empty() { None } else { Some(reasoning.clone()) }, + tool_calls: Some(tc_requests), + tool_call_id: None, + }); + + let mut current_batch_node_ids: Vec = Vec::new(); + + for tc in &tool_calls { + total_tool_calls += 1; + let node_id = uuid::Uuid::new_v4().to_string(); + + let mut node = AttackChainNode::new( + session_id.clone(), + node_id.clone(), + tc.name.clone(), + tc.arguments.clone(), + reasoning.clone(), + ); + // Link to previous iteration's nodes + node.parent_node_ids = prev_node_ids.clone(); + node.status = AttackNodeStatus::Running; + node.started_at = Some(chrono::Utc::now()); + let _ = self.db.attack_chain_nodes().insert_one(&node).await; + current_batch_node_ids.push(node_id.clone()); + + let _ = self.event_tx.send(PentestEvent::ToolStart { + node_id: node_id.clone(), + tool_name: tc.name.clone(), + input: tc.arguments.clone(), + }); + + let result = if let Some(tool) = self.tool_registry.get(&tc.name) { + match tool.execute(tc.arguments.clone(), &tool_context).await { + Ok(result) => { + total_successes += 1; + let findings_count = result.findings.len() as u32; + total_findings += findings_count; + + let mut finding_ids: Vec = Vec::new(); + for mut finding in result.findings { + finding.scan_run_id = session_id.clone(); + finding.session_id = Some(session_id.clone()); + let insert_result = + self.db.dast_findings().insert_one(&finding).await; + if let Ok(res) = &insert_result { + finding_ids.push(res.inserted_id.as_object_id().map(|oid| oid.to_hex()).unwrap_or_default()); + } + let _ = + self.event_tx.send(PentestEvent::Finding { + finding_id: finding + .id + .map(|oid| oid.to_hex()) + .unwrap_or_default(), + title: finding.title.clone(), + severity: finding.severity.to_string(), + }); + } + + // Compute risk score based on findings severity + let risk_score: Option = if findings_count > 0 { + Some(std::cmp::min( + 100, + (findings_count as u8).saturating_mul(15).saturating_add(20), + )) + } else { + None + }; + + let _ = self.event_tx.send(PentestEvent::ToolComplete { + node_id: node_id.clone(), + summary: result.summary.clone(), + findings_count, + }); + + let finding_ids_bson: Vec = finding_ids + .iter() + .map(|id| mongodb::bson::Bson::String(id.clone())) + .collect(); + + let mut update_doc = doc! { + "status": "completed", + "tool_output": mongodb::bson::to_bson(&result.data) + .unwrap_or(mongodb::bson::Bson::Null), + "completed_at": mongodb::bson::DateTime::now(), + "findings_produced": finding_ids_bson, + }; + if let Some(rs) = risk_score { + update_doc.insert("risk_score", rs as i32); + } + + let _ = self + .db + .attack_chain_nodes() + .update_one( + doc! { + "session_id": &session_id, + "node_id": &node_id, + }, + doc! { "$set": update_doc }, + ) + .await; + + serde_json::json!({ + "summary": result.summary, + "findings_count": findings_count, + "data": result.data, + }) + .to_string() + } + Err(e) => { + let _ = self + .db + .attack_chain_nodes() + .update_one( + doc! { + "session_id": &session_id, + "node_id": &node_id, + }, + doc! { "$set": { + "status": "failed", + "completed_at": mongodb::bson::DateTime::now(), + }}, + ) + .await; + format!("Tool execution failed: {e}") + } + } + } else { + format!("Unknown tool: {}", tc.name) + }; + + messages.push(ChatMessage { + role: "tool".to_string(), + content: Some(result), + tool_calls: None, + tool_call_id: Some(tc.id.clone()), + }); + } + + // Advance parent links so next iteration's nodes connect to this batch + prev_node_ids = current_batch_node_ids; + + if let Some(sid) = session.id { + let _ = self + .db + .pentest_sessions() + .update_one( + doc! { "_id": sid }, + doc! { "$set": { + "tool_invocations": total_tool_calls as i64, + "tool_successes": total_successes as i64, + "findings_count": total_findings as i64, + }}, + ) + .await; + } + } + } + } + + if let Some(sid) = session.id { + let _ = self + .db + .pentest_sessions() + .update_one( + doc! { "_id": sid }, + doc! { "$set": { + "status": "completed", + "completed_at": mongodb::bson::DateTime::now(), + "tool_invocations": total_tool_calls as i64, + "tool_successes": total_successes as i64, + "findings_count": total_findings as i64, + }}, + ) + .await; + } + + let _ = self.event_tx.send(PentestEvent::Complete { + summary: format!( + "Pentest complete. {} findings from {} tool invocations.", + total_findings, total_tool_calls + ), + }); + + Ok(()) + } + + // ── Code-Awareness: Gather context from linked repo ───────── + + /// Fetch SAST findings, SBOM entries (with CVEs), and code graph entry points + /// for the repo linked to this DAST target. + async fn gather_repo_context( + &self, + target: &DastTarget, + ) -> (Vec, Vec, Vec) { + let Some(repo_id) = &target.repo_id else { + return (Vec::new(), Vec::new(), Vec::new()); + }; + + let sast_findings = self.fetch_sast_findings(repo_id).await; + let sbom_entries = self.fetch_vulnerable_sbom(repo_id).await; + let code_context = self.fetch_code_context(repo_id, &sast_findings).await; + + tracing::info!( + repo_id, + sast_findings = sast_findings.len(), + vulnerable_deps = sbom_entries.len(), + code_hints = code_context.len(), + "Gathered code-awareness context for pentest" + ); + + (sast_findings, sbom_entries, code_context) + } + + /// Fetch open/triaged SAST findings for the repo (not false positives or resolved) + async fn fetch_sast_findings(&self, repo_id: &str) -> Vec { + let cursor = self + .db + .findings() + .find(doc! { + "repo_id": repo_id, + "status": { "$in": ["open", "triaged"] }, + }) + .sort(doc! { "severity": -1 }) + .limit(100) + .await; + + match cursor { + Ok(mut c) => { + let mut results = Vec::new(); + while let Some(Ok(f)) = c.next().await { + results.push(f); + } + results + } + Err(e) => { + tracing::warn!("Failed to fetch SAST findings for pentest: {e}"); + Vec::new() + } + } + } + + /// Fetch SBOM entries that have known vulnerabilities + async fn fetch_vulnerable_sbom(&self, repo_id: &str) -> Vec { + let cursor = self + .db + .sbom_entries() + .find(doc! { + "repo_id": repo_id, + "known_vulnerabilities": { "$exists": true, "$ne": [] }, + }) + .limit(50) + .await; + + match cursor { + Ok(mut c) => { + let mut results = Vec::new(); + while let Some(Ok(e)) = c.next().await { + results.push(e); + } + results + } + Err(e) => { + tracing::warn!("Failed to fetch vulnerable SBOM entries: {e}"); + Vec::new() + } + } + } + + /// Build CodeContextHint objects from the code knowledge graph. + /// Maps entry points to their source files and links SAST findings. + async fn fetch_code_context( + &self, + repo_id: &str, + sast_findings: &[Finding], + ) -> Vec { + // Get entry point nodes from the code graph + let cursor = self + .db + .graph_nodes() + .find(doc! { + "repo_id": repo_id, + "is_entry_point": true, + }) + .limit(50) + .await; + + let nodes = match cursor { + Ok(mut c) => { + let mut results = Vec::new(); + while let Some(Ok(n)) = c.next().await { + results.push(n); + } + results + } + Err(_) => return Vec::new(), + }; + + // Build hints by matching graph nodes to SAST findings by file path + nodes + .into_iter() + .map(|node| { + // Find SAST findings in the same file + let linked_vulns: Vec = sast_findings + .iter() + .filter(|f| { + f.file_path.as_deref() == Some(&node.file_path) + }) + .map(|f| { + format!( + "[{}] {}: {} (line {})", + f.severity, + f.scanner, + f.title, + f.line_number.unwrap_or(0) + ) + }) + .collect(); + + CodeContextHint { + endpoint_pattern: node.qualified_name.clone(), + handler_function: node.name.clone(), + file_path: node.file_path.clone(), + code_snippet: String::new(), // Could fetch from embeddings + known_vulnerabilities: linked_vulns, + } + }) + .collect() + } + + // ── System Prompt Builder ─────────────────────────────────── + + async fn build_system_prompt( + &self, + session: &PentestSession, + target: &DastTarget, + sast_findings: &[Finding], + sbom_entries: &[SbomEntry], + code_context: &[CodeContextHint], + ) -> String { + let tool_names = self.tool_registry.list_names().join(", "); + let strategy_guidance = match session.strategy { + PentestStrategy::Quick => { + "Focus on the most common and impactful vulnerabilities. Run a quick recon, then target the highest-risk areas." + } + PentestStrategy::Comprehensive => { + "Perform a thorough assessment covering all vulnerability types. Start with recon, then systematically test each attack surface." + } + PentestStrategy::Targeted => { + "Focus specifically on areas highlighted by SAST findings and known CVEs. Prioritize exploiting known weaknesses." + } + PentestStrategy::Aggressive => { + "Use all available tools aggressively. Test with maximum payloads and attempt full exploitation." + } + PentestStrategy::Stealth => { + "Minimize noise. Use fewer requests, avoid aggressive payloads. Focus on passive analysis and targeted probes." + } + }; + + // Build SAST findings section + let sast_section = if sast_findings.is_empty() { + String::from("No SAST findings available for this target.") + } else { + let critical = sast_findings + .iter() + .filter(|f| f.severity == Severity::Critical) + .count(); + let high = sast_findings + .iter() + .filter(|f| f.severity == Severity::High) + .count(); + + let mut section = format!( + "{} open findings ({} critical, {} high):\n", + sast_findings.len(), + critical, + high + ); + + // List the most important findings (critical/high first, up to 20) + for f in sast_findings.iter().take(20) { + let file_info = f + .file_path + .as_ref() + .map(|p| { + format!( + " in {}:{}", + p, + f.line_number.unwrap_or(0) + ) + }) + .unwrap_or_default(); + let status_note = match f.status { + FindingStatus::Triaged => " [TRIAGED]", + _ => "", + }; + section.push_str(&format!( + "- [{sev}] {title}{file}{status}\n", + sev = f.severity, + title = f.title, + file = file_info, + status = status_note, + )); + if let Some(cwe) = &f.cwe { + section.push_str(&format!(" CWE: {cwe}\n")); + } + } + if sast_findings.len() > 20 { + section.push_str(&format!( + "... and {} more findings\n", + sast_findings.len() - 20 + )); + } + section + }; + + // Build SBOM/CVE section + let sbom_section = if sbom_entries.is_empty() { + String::from("No vulnerable dependencies identified.") + } else { + let mut section = format!( + "{} dependencies with known vulnerabilities:\n", + sbom_entries.len() + ); + for entry in sbom_entries.iter().take(15) { + let cve_ids: Vec<&str> = entry + .known_vulnerabilities + .iter() + .map(|v| v.id.as_str()) + .collect(); + section.push_str(&format!( + "- {} {} ({}): {}\n", + entry.name, + entry.version, + entry.package_manager, + cve_ids.join(", ") + )); + } + if sbom_entries.len() > 15 { + section.push_str(&format!( + "... and {} more vulnerable dependencies\n", + sbom_entries.len() - 15 + )); + } + section + }; + + // Build code context section + let code_section = if code_context.is_empty() { + String::from("No code knowledge graph available for this target.") + } else { + let with_vulns = code_context + .iter() + .filter(|c| !c.known_vulnerabilities.is_empty()) + .count(); + + let mut section = format!( + "{} entry points identified ({} with linked SAST findings):\n", + code_context.len(), + with_vulns + ); + + for hint in code_context.iter().take(20) { + section.push_str(&format!( + "- {} ({})\n", + hint.endpoint_pattern, hint.file_path + )); + for vuln in &hint.known_vulnerabilities { + section.push_str(&format!(" SAST: {vuln}\n")); + } + } + section + }; + + format!( + r#"You are an expert penetration tester conducting an authorized security assessment. + +## Target +- **Name**: {target_name} +- **URL**: {base_url} +- **Type**: {target_type} +- **Rate Limit**: {rate_limit} req/s +- **Destructive Tests Allowed**: {allow_destructive} +- **Linked Repository**: {repo_linked} + +## Strategy +{strategy_guidance} + +## SAST Findings (Static Analysis) +{sast_section} + +## Vulnerable Dependencies (SBOM) +{sbom_section} + +## Code Entry Points (Knowledge Graph) +{code_section} + +## Available Tools +{tool_names} + +## Instructions +1. Start by running reconnaissance (recon tool) to fingerprint the target and discover technologies. +2. Run the OpenAPI parser to discover API endpoints from specs. +3. Check infrastructure: DNS, DMARC, TLS, security headers, cookies, CSP, CORS. +4. Based on SAST findings, prioritize testing endpoints where vulnerabilities were found in code. +5. For each vulnerability type found in SAST, use the corresponding DAST tool to verify exploitability. +6. If vulnerable dependencies are listed, try to trigger known CVE conditions against the running application. +7. Test rate limiting on critical endpoints (login, API). +8. Check for console.log leakage in frontend JavaScript. +9. Analyze tool results and chain findings — if one vulnerability enables others, explore the chain. +10. When testing is complete, provide a structured summary with severity and remediation. +11. Always explain your reasoning before invoking each tool. +12. When done, say "Testing complete" followed by a final summary. + +## Important +- This is an authorized penetration test. All testing is permitted within the target scope. +- Respect the rate limit of {rate_limit} requests per second. +- Only use destructive tests if explicitly allowed ({allow_destructive}). +- Use SAST findings to guide your testing — they tell you WHERE in the code vulnerabilities exist. +- Use SBOM data to understand what technologies and versions the target runs. +"#, + target_name = target.name, + base_url = target.base_url, + target_type = target.target_type, + rate_limit = target.rate_limit, + allow_destructive = target.allow_destructive, + repo_linked = target.repo_id.as_deref().unwrap_or("None"), + ) + } +} diff --git a/compliance-agent/src/pentest/report.rs b/compliance-agent/src/pentest/report.rs new file mode 100644 index 0000000..6997ded --- /dev/null +++ b/compliance-agent/src/pentest/report.rs @@ -0,0 +1,1601 @@ +use std::io::{Cursor, Write}; + +use compliance_core::models::dast::DastFinding; +use compliance_core::models::pentest::{AttackChainNode, PentestSession}; +use sha2::{Digest, Sha256}; +use zip::write::SimpleFileOptions; +use zip::AesMode; + +/// Report archive with metadata +pub struct ReportArchive { + /// The password-protected ZIP bytes + pub archive: Vec, + /// SHA-256 hex digest of the archive + pub sha256: String, +} + +/// Report context gathered from the database +pub struct ReportContext { + pub session: PentestSession, + pub target_name: String, + pub target_url: String, + pub findings: Vec, + pub attack_chain: Vec, + pub requester_name: String, + pub requester_email: String, +} + +/// Generate a password-protected ZIP archive containing the pentest report. +/// +/// The archive contains: +/// - `report.pdf` — Professional pentest report (PDF) +/// - `report.html` — HTML source (fallback) +/// - `findings.json` — Raw findings data +/// - `attack-chain.json` — Attack chain timeline +/// +/// Files are encrypted with AES-256 inside the ZIP (standard WinZip AES format, +/// supported by 7-Zip, WinRAR, macOS Archive Utility, etc.). +pub async fn generate_encrypted_report( + ctx: &ReportContext, + password: &str, +) -> Result { + let html = build_html_report(ctx); + + // Convert HTML to PDF via headless Chrome + let pdf_bytes = html_to_pdf(&html).await?; + + let zip_bytes = build_zip(ctx, password, &html, &pdf_bytes) + .map_err(|e| format!("Failed to create archive: {e}"))?; + + let mut hasher = Sha256::new(); + hasher.update(&zip_bytes); + let sha256 = hex::encode(hasher.finalize()); + + Ok(ReportArchive { archive: zip_bytes, sha256 }) +} + +/// Convert HTML string to PDF bytes using headless Chrome/Chromium. +async fn html_to_pdf(html: &str) -> Result, String> { + let tmp_dir = std::env::temp_dir(); + let run_id = uuid::Uuid::new_v4().to_string(); + let html_path = tmp_dir.join(format!("pentest-report-{run_id}.html")); + let pdf_path = tmp_dir.join(format!("pentest-report-{run_id}.pdf")); + + // Write HTML to temp file + std::fs::write(&html_path, html) + .map_err(|e| format!("Failed to write temp HTML: {e}"))?; + + // Find Chrome/Chromium binary + let chrome_bin = find_chrome_binary() + .ok_or_else(|| "Chrome/Chromium not found. Install google-chrome or chromium to generate PDF reports.".to_string())?; + + tracing::info!(chrome = %chrome_bin, "Generating PDF report via headless Chrome"); + + let html_url = format!("file://{}", html_path.display()); + + let output = tokio::process::Command::new(&chrome_bin) + .args([ + "--headless", + "--disable-gpu", + "--no-sandbox", + "--disable-software-rasterizer", + "--run-all-compositor-stages-before-draw", + "--disable-dev-shm-usage", + &format!("--print-to-pdf={}", pdf_path.display()), + "--no-pdf-header-footer", + &html_url, + ]) + .output() + .await + .map_err(|e| format!("Failed to run Chrome: {e}"))?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + // Clean up temp files + let _ = std::fs::remove_file(&html_path); + let _ = std::fs::remove_file(&pdf_path); + return Err(format!("Chrome PDF generation failed: {stderr}")); + } + + let pdf_bytes = std::fs::read(&pdf_path) + .map_err(|e| format!("Failed to read generated PDF: {e}"))?; + + // Clean up temp files + let _ = std::fs::remove_file(&html_path); + let _ = std::fs::remove_file(&pdf_path); + + if pdf_bytes.is_empty() { + return Err("Chrome produced an empty PDF".to_string()); + } + + tracing::info!(size_kb = pdf_bytes.len() / 1024, "PDF report generated"); + Ok(pdf_bytes) +} + +/// Search for Chrome/Chromium binary on the system. +fn find_chrome_binary() -> Option { + let candidates = [ + "google-chrome-stable", + "google-chrome", + "chromium-browser", + "chromium", + ]; + for name in &candidates { + if let Ok(output) = std::process::Command::new("which").arg(name).output() { + if output.status.success() { + let path = String::from_utf8_lossy(&output.stdout).trim().to_string(); + if !path.is_empty() { + return Some(path); + } + } + } + } + None +} + +fn build_zip( + ctx: &ReportContext, + password: &str, + html: &str, + pdf: &[u8], +) -> Result, zip::result::ZipError> { + let buf = Cursor::new(Vec::new()); + let mut zip = zip::ZipWriter::new(buf); + + let options = SimpleFileOptions::default() + .compression_method(zip::CompressionMethod::Deflated) + .with_aes_encryption(AesMode::Aes256, password); + + // report.pdf (primary) + zip.start_file("report.pdf", options.clone())?; + zip.write_all(pdf)?; + + // report.html (fallback) + zip.start_file("report.html", options.clone())?; + zip.write_all(html.as_bytes())?; + + // findings.json + let findings_json = + serde_json::to_string_pretty(&ctx.findings).unwrap_or_else(|_| "[]".to_string()); + zip.start_file("findings.json", options.clone())?; + zip.write_all(findings_json.as_bytes())?; + + // attack-chain.json + let chain_json = + serde_json::to_string_pretty(&ctx.attack_chain).unwrap_or_else(|_| "[]".to_string()); + zip.start_file("attack-chain.json", options)?; + zip.write_all(chain_json.as_bytes())?; + + let cursor = zip.finish()?; + Ok(cursor.into_inner()) +} + +fn build_html_report(ctx: &ReportContext) -> String { + let session = &ctx.session; + let session_id = session + .id + .map(|oid| oid.to_hex()) + .unwrap_or_else(|| "-".to_string()); + let date_str = session.started_at.format("%B %d, %Y at %H:%M UTC").to_string(); + let date_short = session.started_at.format("%B %d, %Y").to_string(); + let completed_str = session + .completed_at + .map(|d| d.format("%B %d, %Y at %H:%M UTC").to_string()) + .unwrap_or_else(|| "In Progress".to_string()); + + let critical = ctx.findings.iter().filter(|f| f.severity.to_string() == "critical").count(); + let high = ctx.findings.iter().filter(|f| f.severity.to_string() == "high").count(); + let medium = ctx.findings.iter().filter(|f| f.severity.to_string() == "medium").count(); + let low = ctx.findings.iter().filter(|f| f.severity.to_string() == "low").count(); + let info = ctx.findings.iter().filter(|f| f.severity.to_string() == "info").count(); + let exploitable = ctx.findings.iter().filter(|f| f.exploitable).count(); + let total = ctx.findings.len(); + + let overall_risk = if critical > 0 { + "CRITICAL" + } else if high > 0 { + "HIGH" + } else if medium > 0 { + "MEDIUM" + } else if low > 0 { + "LOW" + } else { + "INFORMATIONAL" + }; + + let risk_color = match overall_risk { + "CRITICAL" => "#991b1b", + "HIGH" => "#c2410c", + "MEDIUM" => "#a16207", + "LOW" => "#1d4ed8", + _ => "#4b5563", + }; + + // Risk score 0-100 + let risk_score: usize = std::cmp::min( + 100, + critical * 25 + high * 15 + medium * 8 + low * 3 + info * 1, + ); + + // Collect unique tool names used + let tool_names: Vec = { + let mut names: Vec = ctx + .attack_chain + .iter() + .map(|n| n.tool_name.clone()) + .collect(); + names.sort(); + names.dedup(); + names + }; + + // Severity distribution bar + let severity_bar = if total > 0 { + let crit_pct = (critical as f64 / total as f64 * 100.0) as usize; + let high_pct = (high as f64 / total as f64 * 100.0) as usize; + let med_pct = (medium as f64 / total as f64 * 100.0) as usize; + let low_pct = (low as f64 / total as f64 * 100.0) as usize; + let info_pct = 100_usize.saturating_sub(crit_pct + high_pct + med_pct + low_pct); + + let mut bar = String::from(r#"
"#); + if critical > 0 { + bar.push_str(&format!( + r#"
{}
"#, + std::cmp::max(crit_pct, 4), critical + )); + } + if high > 0 { + bar.push_str(&format!( + r#"
{}
"#, + std::cmp::max(high_pct, 4), high + )); + } + if medium > 0 { + bar.push_str(&format!( + r#"
{}
"#, + std::cmp::max(med_pct, 4), medium + )); + } + if low > 0 { + bar.push_str(&format!( + r#"
{}
"#, + std::cmp::max(low_pct, 4), low + )); + } + if info > 0 { + bar.push_str(&format!( + r#"
{}
"#, + std::cmp::max(info_pct, 4), info + )); + } + bar.push_str("
"); + bar.push_str(r#"
"#); + if critical > 0 { bar.push_str(r#" Critical"#); } + if high > 0 { bar.push_str(r#" High"#); } + if medium > 0 { bar.push_str(r#" Medium"#); } + if low > 0 { bar.push_str(r#" Low"#); } + if info > 0 { bar.push_str(r#" Info"#); } + bar.push_str("
"); + bar + } else { + String::new() + }; + + // Build findings grouped by severity + let severity_order = ["critical", "high", "medium", "low", "info"]; + let severity_labels = ["Critical", "High", "Medium", "Low", "Informational"]; + let severity_colors = ["#991b1b", "#c2410c", "#a16207", "#1d4ed8", "#4b5563"]; + + let mut findings_html = String::new(); + let mut finding_num = 0usize; + + for (si, &sev_key) in severity_order.iter().enumerate() { + let sev_findings: Vec<&DastFinding> = ctx + .findings + .iter() + .filter(|f| f.severity.to_string() == sev_key) + .collect(); + if sev_findings.is_empty() { + continue; + } + + findings_html.push_str(&format!( + r#"

{label} ({count})

"#, + color = severity_colors[si], + label = severity_labels[si], + count = sev_findings.len(), + )); + + for f in sev_findings { + finding_num += 1; + let sev_color = severity_colors[si]; + let exploitable_badge = if f.exploitable { + r#"EXPLOITABLE"# + } else { + "" + }; + let cwe_cell = f + .cwe + .as_deref() + .map(|c| format!("CWE{}", html_escape(c))) + .unwrap_or_default(); + let param_row = f + .parameter + .as_deref() + .map(|p| format!("Parameter{}", html_escape(p))) + .unwrap_or_default(); + let remediation = f + .remediation + .as_deref() + .unwrap_or("Refer to industry best practices for this vulnerability class."); + + let evidence_html = if f.evidence.is_empty() { + String::new() + } else { + let mut eh = String::from(r#"
Evidence
"#); + for ev in &f.evidence { + let payload_info = ev + .payload + .as_deref() + .map(|p| format!("
Payload: {}", html_escape(p))) + .unwrap_or_default(); + eh.push_str(&format!( + "", + html_escape(&ev.request_method), + html_escape(&ev.request_url), + ev.response_status, + ev.response_snippet + .as_deref() + .map(|s| html_escape(s)) + .unwrap_or_default(), + payload_info, + )); + } + eh.push_str("
RequestStatusDetails
{} {}{}{}{}
"); + eh + }; + + let linked_sast = f + .linked_sast_finding_id + .as_deref() + .map(|id| { + format!( + r#"
Correlated SAST Finding: {id}
"# + ) + }) + .unwrap_or_default(); + + findings_html.push_str(&format!( + r#" +
+
+ F-{num:03} + {title} + {exploitable_badge} +
+ + + + {param_row} + {cwe_cell} +
Type{vuln_type}
Endpoint{method} {endpoint}
+
{description}
+ {evidence_html} + {linked_sast} +
+
Recommendation
+ {remediation} +
+
+ "#, + num = finding_num, + title = html_escape(&f.title), + vuln_type = f.vuln_type, + method = f.method, + endpoint = html_escape(&f.endpoint), + description = html_escape(&f.description), + )); + } + } + + // Build attack chain — group by phase using BFS + let mut chain_html = String::new(); + if !ctx.attack_chain.is_empty() { + // Compute phases via BFS from root nodes + let mut phase_map: std::collections::HashMap = std::collections::HashMap::new(); + let mut queue: std::collections::VecDeque = std::collections::VecDeque::new(); + + for node in &ctx.attack_chain { + if node.parent_node_ids.is_empty() { + let nid = node.node_id.clone(); + if !nid.is_empty() { + phase_map.insert(nid.clone(), 0); + queue.push_back(nid); + } + } + } + + while let Some(nid) = queue.pop_front() { + let parent_phase = phase_map.get(&nid).copied().unwrap_or(0); + for node in &ctx.attack_chain { + if node.parent_node_ids.contains(&nid) { + let child_id = node.node_id.clone(); + if !child_id.is_empty() && !phase_map.contains_key(&child_id) { + phase_map.insert(child_id.clone(), parent_phase + 1); + queue.push_back(child_id); + } + } + } + } + + // Assign phase 0 to any unassigned nodes + for node in &ctx.attack_chain { + let nid = node.node_id.clone(); + if !nid.is_empty() && !phase_map.contains_key(&nid) { + phase_map.insert(nid, 0); + } + } + + // Group nodes by phase + let max_phase = phase_map.values().copied().max().unwrap_or(0); + let phase_labels = ["Reconnaissance", "Enumeration", "Exploitation", "Validation", "Post-Exploitation"]; + + for phase_idx in 0..=max_phase { + let phase_nodes: Vec<&AttackChainNode> = ctx + .attack_chain + .iter() + .filter(|n| { + let nid = n.node_id.clone(); + phase_map.get(&nid).copied().unwrap_or(0) == phase_idx + }) + .collect(); + + if phase_nodes.is_empty() { + continue; + } + + let label = if phase_idx < phase_labels.len() { + phase_labels[phase_idx] + } else { + "Additional Testing" + }; + + chain_html.push_str(&format!( + r#"
+
+ Phase {} + {} + {} step{} +
+
"#, + phase_idx + 1, + label, + phase_nodes.len(), + if phase_nodes.len() == 1 { "" } else { "s" }, + )); + + for (i, node) in phase_nodes.iter().enumerate() { + let status_label = format!("{:?}", node.status); + let status_class = match status_label.to_lowercase().as_str() { + "completed" => "step-completed", + "failed" => "step-failed", + _ => "step-running", + }; + let findings_badge = if !node.findings_produced.is_empty() { + format!( + r#"{} finding{}"#, + node.findings_produced.len(), + if node.findings_produced.len() == 1 { "" } else { "s" }, + ) + } else { + String::new() + }; + let risk_badge = node.risk_score.map(|r| { + let risk_class = if r >= 70 { "risk-high" } else if r >= 40 { "risk-med" } else { "risk-low" }; + format!(r#"Risk: {r}"#) + }).unwrap_or_default(); + + let reasoning_html = if node.llm_reasoning.is_empty() { + String::new() + } else { + format!( + r#"
{}
"#, + html_escape(&node.llm_reasoning) + ) + }; + + chain_html.push_str(&format!( + r#"
+
{num}
+
+
+
+ {tool_name} + {status_label} + {findings_badge} + {risk_badge} +
+ {reasoning_html} +
+
"#, + num = i + 1, + tool_name = html_escape(&node.tool_name), + )); + } + + chain_html.push_str("
"); + } + } + + // Tools methodology table + let tools_table: String = tool_names + .iter() + .enumerate() + .map(|(i, t)| { + let category = tool_category(t); + format!( + "{}{}{}", + i + 1, + html_escape(t), + category, + ) + }) + .collect::>() + .join("\n"); + + // Table of contents + let toc_findings_sub = if !ctx.findings.is_empty() { + let mut sub = String::new(); + let mut fnum = 0usize; + for (si, &sev_key) in severity_order.iter().enumerate() { + let count = ctx.findings.iter().filter(|f| f.severity.to_string() == sev_key).count(); + if count == 0 { continue; } + for f in ctx.findings.iter().filter(|f| f.severity.to_string() == sev_key) { + fnum += 1; + sub.push_str(&format!( + r#"
F-{:03} — {}
"#, + fnum, + html_escape(&f.title), + )); + } + } + sub + } else { + String::new() + }; + + format!( + r##" + + + + +Penetration Test Report — {target_name} + + + + + + + + +
+ + + + + + + + + + + + + + +
CONFIDENTIAL
+ +
Penetration Test Report
+
{target_name}
+ +
+ +
+ Report ID: {session_id}
+ Date: {date_short}
+ Target: {target_url}
+ Prepared for: {requester_name} ({requester_email}) +
+ +
+ Compliance Scanner — AI-Powered Security Assessment Platform +
+
+ + +
+ +
+

Table of Contents

+
1Executive Summary
+
2Scope & Methodology
+
3Findings ({total_findings})
+ {toc_findings_sub} +
4Attack Chain Timeline
+
5Appendix
+
+ + +

1. Executive Summary

+ +
+
+
+
+
+
{risk_score} / 100
+
+
+
Overall Risk: {overall_risk}
+
+ Based on {total_findings} finding{findings_plural} identified across the target application. +
+
+
+ +
+
+
{total_findings}
+
Total Findings
+
+
+
{critical_high}
+
Critical / High
+
+
+
{exploitable_count}
+
Exploitable
+
+
+
{tool_count}
+
Tools Used
+
+
+ +

Severity Distribution

+{severity_bar} + +

+ This report presents the results of an automated penetration test conducted against + {target_name} ({target_url}) using the Compliance Scanner + AI-powered testing engine. A total of {total_findings} vulnerabilities were + identified, of which {exploitable_count} were confirmed exploitable with + working proof-of-concept payloads. The assessment employed {tool_count} security tools + across {tool_invocations} invocations ({success_rate:.0}% success rate). +

+ + +
+

2. Scope & Methodology

+ +

+ The assessment was performed using an AI-driven orchestrator that autonomously selects and + executes security testing tools based on the target's attack surface, technology stack, and + any available static analysis (SAST) findings and SBOM data. +

+ +

Engagement Details

+ + + + + + + + +
Target{target_name}
URL{target_url}
Strategy{strategy}
Status{status}
Started{date_str}
Completed{completed_str}
Tool Invocations{tool_invocations} ({tool_successes} successful, {success_rate:.1}% success rate)
+ +

Tools Employed

+ + + {tools_table} +
#ToolCategory
+ + +
+

3. Findings

+ +{findings_section} + + +
+

4. Attack Chain Timeline

+ +

+ The following sequence shows each tool invocation made by the AI orchestrator during the assessment, + grouped by phase. Each step includes the tool's name, execution status, and the AI's reasoning + for choosing that action. +

+ +
+ {chain_section} +
+ + +
+

5. Appendix

+ +

Severity Definitions

+ + + + + + +
CriticalVulnerabilities that can be exploited remotely without authentication to execute arbitrary code, exfiltrate sensitive data, or fully compromise the system.
HighVulnerabilities that allow significant unauthorized access or data exposure, typically requiring minimal user interaction or privileges.
MediumVulnerabilities that may lead to limited data exposure or require specific conditions to exploit, but still represent meaningful risk.
LowMinor issues with limited direct impact. May contribute to broader attack chains or indicate defense-in-depth weaknesses.
InfoObservations and best-practice recommendations that do not represent direct security vulnerabilities.
+ +

Disclaimer

+

+ This report was generated by an automated AI-powered penetration testing engine. While the system + employs advanced techniques to identify vulnerabilities, no automated assessment can guarantee + complete coverage. The results should be reviewed by qualified security professionals and validated + in the context of the target application's threat model. Findings are point-in-time observations + and may change as the application evolves. +

+ + +
+
Compliance Scanner
+
AI-Powered Security Assessment Platform
+
This document is confidential and intended solely for the named recipient.
+
Report ID: {session_id}
+
+ +
+ + +"##, + target_name = html_escape(&ctx.target_name), + target_url = html_escape(&ctx.target_url), + session_id = html_escape(&session_id), + date_str = date_str, + date_short = date_short, + completed_str = completed_str, + requester_name = html_escape(&ctx.requester_name), + requester_email = html_escape(&ctx.requester_email), + risk_color = risk_color, + risk_score = risk_score, + overall_risk = overall_risk, + total_findings = total, + findings_plural = if total == 1 { "" } else { "s" }, + critical_high = format!("{} / {}", critical, high), + exploitable_count = exploitable, + tool_count = tool_names.len(), + strategy = session.strategy, + status = session.status, + tool_invocations = session.tool_invocations, + tool_successes = session.tool_successes, + success_rate = session.success_rate(), + severity_bar = severity_bar, + tools_table = tools_table, + toc_findings_sub = toc_findings_sub, + findings_section = if ctx.findings.is_empty() { + "

No vulnerabilities were identified during this assessment.

".to_string() + } else { + findings_html + }, + chain_section = if ctx.attack_chain.is_empty() { + "

No attack chain steps recorded.

".to_string() + } else { + chain_html + }, + ) +} + +fn tool_category(tool_name: &str) -> &'static str { + let name = tool_name.to_lowercase(); + if name.contains("nmap") || name.contains("port") { return "Network Reconnaissance"; } + if name.contains("nikto") || name.contains("header") { return "Web Server Analysis"; } + if name.contains("zap") || name.contains("spider") || name.contains("crawl") { return "Web Application Scanning"; } + if name.contains("sqlmap") || name.contains("sqli") || name.contains("sql") { return "SQL Injection Testing"; } + if name.contains("xss") || name.contains("cross-site") { return "Cross-Site Scripting Testing"; } + if name.contains("dir") || name.contains("brute") || name.contains("fuzz") || name.contains("gobuster") { return "Directory Enumeration"; } + if name.contains("ssl") || name.contains("tls") || name.contains("cert") { return "SSL/TLS Analysis"; } + if name.contains("api") || name.contains("endpoint") { return "API Security Testing"; } + if name.contains("auth") || name.contains("login") || name.contains("credential") { return "Authentication Testing"; } + if name.contains("cors") { return "CORS Testing"; } + if name.contains("csrf") { return "CSRF Testing"; } + if name.contains("nuclei") || name.contains("template") { return "Vulnerability Scanning"; } + if name.contains("whatweb") || name.contains("tech") || name.contains("wappalyzer") { return "Technology Fingerprinting"; } + "Security Testing" +} + +fn html_escape(s: &str) -> String { + s.replace('&', "&") + .replace('<', "<") + .replace('>', ">") + .replace('"', """) +} diff --git a/compliance-core/src/models/dast.rs b/compliance-core/src/models/dast.rs index 7ff9599..4972aca 100644 --- a/compliance-core/src/models/dast.rs +++ b/compliance-core/src/models/dast.rs @@ -176,6 +176,16 @@ pub enum DastVulnType { InformationDisclosure, SecurityMisconfiguration, BrokenAuth, + DnsMisconfiguration, + EmailSecurity, + TlsMisconfiguration, + CookieSecurity, + CspIssue, + CorsMisconfiguration, + RateLimitAbsent, + ConsoleLogLeakage, + SecurityHeaderMissing, + KnownCveExploit, Other, } @@ -192,6 +202,16 @@ impl std::fmt::Display for DastVulnType { Self::InformationDisclosure => write!(f, "information_disclosure"), Self::SecurityMisconfiguration => write!(f, "security_misconfiguration"), Self::BrokenAuth => write!(f, "broken_auth"), + Self::DnsMisconfiguration => write!(f, "dns_misconfiguration"), + Self::EmailSecurity => write!(f, "email_security"), + Self::TlsMisconfiguration => write!(f, "tls_misconfiguration"), + Self::CookieSecurity => write!(f, "cookie_security"), + Self::CspIssue => write!(f, "csp_issue"), + Self::CorsMisconfiguration => write!(f, "cors_misconfiguration"), + Self::RateLimitAbsent => write!(f, "rate_limit_absent"), + Self::ConsoleLogLeakage => write!(f, "console_log_leakage"), + Self::SecurityHeaderMissing => write!(f, "security_header_missing"), + Self::KnownCveExploit => write!(f, "known_cve_exploit"), Self::Other => write!(f, "other"), } } @@ -244,6 +264,8 @@ pub struct DastFinding { pub remediation: Option, /// Linked SAST finding ID (if correlated) pub linked_sast_finding_id: Option, + /// Pentest session that produced this finding (if AI-driven) + pub session_id: Option, #[serde(with = "super::serde_helpers::bson_datetime")] pub created_at: DateTime, } @@ -276,6 +298,7 @@ impl DastFinding { evidence: Vec::new(), remediation: None, linked_sast_finding_id: None, + session_id: None, created_at: Utc::now(), } } diff --git a/compliance-core/src/models/mod.rs b/compliance-core/src/models/mod.rs index 8d9f064..daf0503 100644 --- a/compliance-core/src/models/mod.rs +++ b/compliance-core/src/models/mod.rs @@ -7,6 +7,7 @@ pub mod finding; pub mod graph; pub mod issue; pub mod mcp; +pub mod pentest; pub mod repository; pub mod sbom; pub mod scan; @@ -26,6 +27,11 @@ pub use graph::{ }; pub use issue::{IssueStatus, TrackerIssue, TrackerType}; pub use mcp::{McpServerConfig, McpServerStatus, McpTransport}; +pub use pentest::{ + AttackChainNode, AttackNodeStatus, CodeContextHint, PentestEvent, PentestMessage, + PentestSession, PentestStats, PentestStatus, PentestStrategy, SeverityDistribution, + ToolCallRecord, +}; pub use repository::{ScanTrigger, TrackedRepository}; pub use sbom::{SbomEntry, VulnRef}; pub use scan::{ScanPhase, ScanRun, ScanRunStatus, ScanType}; diff --git a/compliance-core/src/models/pentest.rs b/compliance-core/src/models/pentest.rs new file mode 100644 index 0000000..dc0fa4d --- /dev/null +++ b/compliance-core/src/models/pentest.rs @@ -0,0 +1,294 @@ +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; + +/// Status of a pentest session +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum PentestStatus { + Running, + Paused, + Completed, + Failed, +} + +impl std::fmt::Display for PentestStatus { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Running => write!(f, "running"), + Self::Paused => write!(f, "paused"), + Self::Completed => write!(f, "completed"), + Self::Failed => write!(f, "failed"), + } + } +} + +/// Strategy for the AI pentest orchestrator +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum PentestStrategy { + /// Quick scan focusing on common vulnerabilities + Quick, + /// Standard comprehensive scan + Comprehensive, + /// Focus on specific vulnerability types guided by SAST/SBOM + Targeted, + /// Aggressive testing with more payloads and deeper exploitation + Aggressive, + /// Stealth mode with slower rate and fewer noisy payloads + Stealth, +} + +impl std::fmt::Display for PentestStrategy { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Quick => write!(f, "quick"), + Self::Comprehensive => write!(f, "comprehensive"), + Self::Targeted => write!(f, "targeted"), + Self::Aggressive => write!(f, "aggressive"), + Self::Stealth => write!(f, "stealth"), + } + } +} + +/// A pentest session initiated via the chat interface +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PentestSession { + #[serde(rename = "_id", skip_serializing_if = "Option::is_none")] + pub id: Option, + pub target_id: String, + /// Linked repository for code-aware testing + pub repo_id: Option, + pub status: PentestStatus, + pub strategy: PentestStrategy, + pub created_by: Option, + /// Total number of tool invocations in this session + pub tool_invocations: u32, + /// Total successful tool invocations + pub tool_successes: u32, + /// Number of findings discovered + pub findings_count: u32, + /// Number of confirmed exploitable findings + pub exploitable_count: u32, + #[serde(with = "super::serde_helpers::bson_datetime")] + pub started_at: DateTime, + #[serde(default, with = "super::serde_helpers::opt_bson_datetime")] + pub completed_at: Option>, +} + +impl PentestSession { + pub fn new(target_id: String, strategy: PentestStrategy) -> Self { + Self { + id: None, + target_id, + repo_id: None, + status: PentestStatus::Running, + strategy, + created_by: None, + tool_invocations: 0, + tool_successes: 0, + findings_count: 0, + exploitable_count: 0, + started_at: Utc::now(), + completed_at: None, + } + } + + pub fn success_rate(&self) -> f64 { + if self.tool_invocations == 0 { + return 100.0; + } + (self.tool_successes as f64 / self.tool_invocations as f64) * 100.0 + } +} + +/// Status of a node in the attack chain +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum AttackNodeStatus { + Pending, + Running, + Completed, + Failed, + Skipped, +} + +/// A single step in the LLM-driven attack chain DAG +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AttackChainNode { + #[serde(rename = "_id", skip_serializing_if = "Option::is_none")] + pub id: Option, + pub session_id: String, + /// Unique ID for DAG references + pub node_id: String, + /// Parent node IDs (multiple for merge nodes) + pub parent_node_ids: Vec, + /// Tool that was invoked + pub tool_name: String, + /// Input parameters passed to the tool + pub tool_input: serde_json::Value, + /// Output from the tool + pub tool_output: Option, + pub status: AttackNodeStatus, + /// LLM's reasoning for choosing this action + pub llm_reasoning: String, + /// IDs of DastFindings produced by this step + pub findings_produced: Vec, + /// Risk score (0-100) assigned by the LLM + pub risk_score: Option, + #[serde(default, with = "super::serde_helpers::opt_bson_datetime")] + pub started_at: Option>, + #[serde(default, with = "super::serde_helpers::opt_bson_datetime")] + pub completed_at: Option>, +} + +impl AttackChainNode { + pub fn new( + session_id: String, + node_id: String, + tool_name: String, + tool_input: serde_json::Value, + llm_reasoning: String, + ) -> Self { + Self { + id: None, + session_id, + node_id, + parent_node_ids: Vec::new(), + tool_name, + tool_input, + tool_output: None, + status: AttackNodeStatus::Pending, + llm_reasoning, + findings_produced: Vec::new(), + risk_score: None, + started_at: None, + completed_at: None, + } + } +} + +/// Chat message within a pentest session +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PentestMessage { + #[serde(rename = "_id", skip_serializing_if = "Option::is_none")] + pub id: Option, + pub session_id: String, + /// "user", "assistant", "tool_result", "system" + pub role: String, + pub content: String, + /// Tool calls made by the assistant in this message + pub tool_calls: Option>, + /// Link to the attack chain node (for tool_result messages) + pub attack_node_id: Option, + #[serde(with = "super::serde_helpers::bson_datetime")] + pub created_at: DateTime, +} + +impl PentestMessage { + pub fn user(session_id: String, content: String) -> Self { + Self { + id: None, + session_id, + role: "user".to_string(), + content, + tool_calls: None, + attack_node_id: None, + created_at: Utc::now(), + } + } + + pub fn assistant(session_id: String, content: String) -> Self { + Self { + id: None, + session_id, + role: "assistant".to_string(), + content, + tool_calls: None, + attack_node_id: None, + created_at: Utc::now(), + } + } + + pub fn tool_result(session_id: String, content: String, node_id: String) -> Self { + Self { + id: None, + session_id, + role: "tool_result".to_string(), + content, + tool_calls: None, + attack_node_id: Some(node_id), + created_at: Utc::now(), + } + } +} + +/// Record of a tool call made by the LLM +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolCallRecord { + pub call_id: String, + pub tool_name: String, + pub arguments: serde_json::Value, + pub result: Option, +} + +/// SSE event types for real-time pentest streaming +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum PentestEvent { + /// LLM is thinking/reasoning + Thinking { reasoning: String }, + /// A tool execution has started + ToolStart { + node_id: String, + tool_name: String, + input: serde_json::Value, + }, + /// A tool execution completed + ToolComplete { + node_id: String, + summary: String, + findings_count: u32, + }, + /// A new finding was discovered + Finding { finding_id: String, title: String, severity: String }, + /// Assistant message (streaming text) + Message { content: String }, + /// Session completed + Complete { summary: String }, + /// Error occurred + Error { message: String }, +} + +/// Aggregated stats for the pentest dashboard +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PentestStats { + pub running_sessions: u32, + pub total_vulnerabilities: u32, + pub total_tool_invocations: u32, + pub tool_success_rate: f64, + pub severity_distribution: SeverityDistribution, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SeverityDistribution { + pub critical: u32, + pub high: u32, + pub medium: u32, + pub low: u32, + pub info: u32, +} + +/// Code context hint linking a discovered endpoint to source code +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CodeContextHint { + /// HTTP route pattern (e.g., "GET /api/users/:id") + pub endpoint_pattern: String, + /// Handler function name + pub handler_function: String, + /// Source file path + pub file_path: String, + /// Relevant code snippet + pub code_snippet: String, + /// SAST findings associated with this code + pub known_vulnerabilities: Vec, +} diff --git a/compliance-core/src/traits/mod.rs b/compliance-core/src/traits/mod.rs index 2d10e8b..4677153 100644 --- a/compliance-core/src/traits/mod.rs +++ b/compliance-core/src/traits/mod.rs @@ -1,9 +1,11 @@ pub mod dast_agent; pub mod graph_builder; pub mod issue_tracker; +pub mod pentest_tool; pub mod scanner; pub use dast_agent::{DastAgent, DastContext, DiscoveredEndpoint, EndpointParameter}; pub use graph_builder::{LanguageParser, ParseOutput}; pub use issue_tracker::IssueTracker; +pub use pentest_tool::{PentestTool, PentestToolContext, PentestToolResult}; pub use scanner::{ScanOutput, Scanner}; diff --git a/compliance-core/src/traits/pentest_tool.rs b/compliance-core/src/traits/pentest_tool.rs new file mode 100644 index 0000000..ad2d73c --- /dev/null +++ b/compliance-core/src/traits/pentest_tool.rs @@ -0,0 +1,63 @@ +use std::future::Future; +use std::pin::Pin; + +use crate::error::CoreError; +use crate::models::dast::{DastFinding, DastTarget}; +use crate::models::finding::Finding; +use crate::models::pentest::CodeContextHint; +use crate::models::sbom::SbomEntry; + +/// Context passed to pentest tools during execution. +/// +/// The HTTP client is not included here because `compliance-core` does not +/// depend on `reqwest`. Tools that need HTTP should hold their own client +/// or receive one via the `compliance-dast` orchestrator. +pub struct PentestToolContext { + /// The DAST target being tested + pub target: DastTarget, + /// Session ID for this pentest run + pub session_id: String, + /// SAST findings for the linked repo (if any) + pub sast_findings: Vec, + /// SBOM entries with known CVEs (if any) + pub sbom_entries: Vec, + /// Code knowledge graph hints mapping endpoints to source code + pub code_context: Vec, + /// Rate limit (requests per second) + pub rate_limit: u32, + /// Whether destructive operations are allowed + pub allow_destructive: bool, +} + +/// Result from a pentest tool execution +pub struct PentestToolResult { + /// Human-readable summary of what the tool found + pub summary: String, + /// DAST findings produced by this tool + pub findings: Vec, + /// Tool-specific structured output data + pub data: serde_json::Value, +} + +/// A tool that the LLM pentest orchestrator can invoke. +/// +/// Each tool represents a specific security testing capability +/// (e.g., SQL injection scanner, DNS checker, TLS analyzer). +/// Uses boxed futures for dyn-compatibility. +pub trait PentestTool: Send + Sync { + /// Tool name for LLM tool_use (e.g., "sql_injection_scanner") + fn name(&self) -> &str; + + /// Human-readable description for the LLM system prompt + fn description(&self) -> &str; + + /// JSON Schema for the tool's input parameters + fn input_schema(&self) -> serde_json::Value; + + /// Execute the tool with the given input + fn execute<'a>( + &'a self, + input: serde_json::Value, + context: &'a PentestToolContext, + ) -> Pin> + Send + 'a>>; +} diff --git a/compliance-dashboard/assets/main.css b/compliance-dashboard/assets/main.css index 8a8cf52..2470dbc 100644 --- a/compliance-dashboard/assets/main.css +++ b/compliance-dashboard/assets/main.css @@ -2767,3 +2767,467 @@ tbody tr:last-child td { .sbom-diff-row-changed { border-left: 3px solid var(--warning); } + +/* ═══════════════════════════════════ + ATTACK CHAIN VISUALIZATION + ═══════════════════════════════════ */ + +/* KPI bar */ +.ac-kpi-bar { + display: flex; + gap: 2px; + margin-bottom: 16px; +} +.ac-kpi-card { + flex: 1; + background: var(--bg-secondary); + border: 1px solid var(--border-color); + padding: 12px 14px; + position: relative; + overflow: hidden; +} +.ac-kpi-card:first-child { border-radius: 10px 0 0 10px; } +.ac-kpi-card:last-child { border-radius: 0 10px 10px 0; } +.ac-kpi-card::before { + content: ''; + position: absolute; + bottom: 0; + left: 0; + right: 0; + height: 2px; +} +.ac-kpi-card:nth-child(1)::before { background: var(--accent, #3b82f6); opacity: 0.4; } +.ac-kpi-card:nth-child(2)::before { background: var(--danger, #dc2626); opacity: 0.5; } +.ac-kpi-card:nth-child(3)::before { background: var(--success, #16a34a); opacity: 0.4; } +.ac-kpi-card:nth-child(4)::before { background: var(--warning, #d97706); opacity: 0.4; } + +.ac-kpi-value { + font-family: var(--font-display); + font-size: 24px; + font-weight: 800; + line-height: 1; + letter-spacing: -0.03em; +} +.ac-kpi-label { + font-family: var(--font-mono, monospace); + font-size: 9px; + color: var(--text-tertiary, #6b7280); + text-transform: uppercase; + letter-spacing: 0.08em; + margin-top: 4px; +} + +/* Phase progress rail */ +.ac-phase-rail { + display: flex; + align-items: flex-start; + margin-bottom: 14px; + position: relative; + padding: 0 8px; +} +.ac-phase-rail::before { + content: ''; + position: absolute; + top: 7px; + left: 8px; + right: 8px; + height: 2px; + background: var(--border-color); + z-index: 0; +} + +.ac-rail-node { + display: flex; + flex-direction: column; + align-items: center; + position: relative; + z-index: 1; + cursor: pointer; + min-width: 56px; + flex: 1; + transition: all 0.15s; +} +.ac-rail-node:hover .ac-rail-dot { transform: scale(1.25); } +.ac-rail-node.active .ac-rail-label { color: var(--accent, #3b82f6); } +.ac-rail-node.active .ac-rail-dot { box-shadow: 0 0 0 3px rgba(59,130,246,0.2), 0 0 12px rgba(59,130,246,0.15); } + +.ac-rail-dot { + width: 14px; + height: 14px; + border-radius: 50%; + border: 2.5px solid var(--bg-primary, #0f172a); + transition: transform 0.2s cubic-bezier(0.16,1,0.3,1); + flex-shrink: 0; +} +.ac-rail-dot.done { background: var(--success, #16a34a); box-shadow: 0 0 8px rgba(22,163,74,0.25); } +.ac-rail-dot.running { background: var(--warning, #d97706); box-shadow: 0 0 10px rgba(217,119,6,0.35); animation: ac-dot-pulse 2s ease-in-out infinite; } +.ac-rail-dot.pending { background: var(--text-tertiary, #6b7280); opacity: 0.5; } +.ac-rail-dot.mixed { background: conic-gradient(var(--success, #16a34a) 0deg 270deg, var(--danger, #dc2626) 270deg 360deg); box-shadow: 0 0 8px rgba(22,163,74,0.2); } + +@keyframes ac-dot-pulse { + 0%, 100% { box-shadow: 0 0 8px rgba(217,119,6,0.35); } + 50% { box-shadow: 0 0 18px rgba(217,119,6,0.55); } +} + +.ac-rail-label { + font-family: var(--font-mono, monospace); + font-size: 9px; + color: var(--text-tertiary, #6b7280); + margin-top: 5px; + letter-spacing: 0.04em; + text-transform: uppercase; + white-space: nowrap; + transition: color 0.15s; +} +.ac-rail-findings { + font-family: var(--font-mono, monospace); + font-size: 9px; + font-weight: 600; + margin-top: 1px; +} +.ac-rail-findings.has { color: var(--danger, #dc2626); } +.ac-rail-findings.none { color: var(--text-tertiary, #6b7280); opacity: 0.4; } + +.ac-rail-heatmap { + display: flex; + gap: 2px; + margin-top: 3px; +} +.ac-hm-cell { + width: 7px; + height: 7px; + border-radius: 1.5px; +} +.ac-hm-cell.ok { background: var(--success, #16a34a); opacity: 0.5; } +.ac-hm-cell.fail { background: var(--danger, #dc2626); opacity: 0.65; } +.ac-hm-cell.run { background: var(--warning, #d97706); opacity: 0.5; animation: ac-pulse 1.5s ease-in-out infinite; } +.ac-hm-cell.wait { background: var(--text-tertiary, #6b7280); opacity: 0.15; } + +.ac-rail-bar { + flex: 1; + height: 2px; + margin-top: 7px; + position: relative; + z-index: 1; +} +.ac-rail-bar-inner { + height: 100%; + border-radius: 1px; +} +.ac-rail-bar-inner.done { background: var(--success, #16a34a); opacity: 0.35; } +.ac-rail-bar-inner.running { background: linear-gradient(to right, var(--success, #16a34a), var(--warning, #d97706)); opacity: 0.35; } + +/* Progress track */ +.ac-progress-track { + height: 3px; + background: var(--border-color); + border-radius: 2px; + overflow: hidden; + margin-bottom: 10px; +} +.ac-progress-fill { + height: 100%; + border-radius: 2px; + background: linear-gradient(90deg, var(--success, #16a34a) 0%, var(--accent, #3b82f6) 100%); + transition: width 0.6s cubic-bezier(0.16,1,0.3,1); +} + +/* Expand all controls */ +.ac-controls { + display: flex; + justify-content: flex-end; + margin-bottom: 6px; +} +.ac-btn-toggle { + font-family: var(--font-body); + font-size: 11px; + color: var(--accent, #3b82f6); + background: none; + border: 1px solid transparent; + cursor: pointer; + padding: 3px 10px; + border-radius: 4px; + transition: all 0.15s; +} +.ac-btn-toggle:hover { + background: rgba(59,130,246,0.08); + border-color: rgba(59,130,246,0.12); +} + +/* Phase accordion */ +.ac-phases { + display: flex; + flex-direction: column; + gap: 2px; +} + +.ac-phase { + animation: ac-phase-in 0.35s cubic-bezier(0.16,1,0.3,1) both; +} +@keyframes ac-phase-in { + from { opacity: 0; transform: translateY(6px); } + to { opacity: 1; transform: translateY(0); } +} + +.ac-phase-header { + display: flex; + align-items: center; + gap: 10px; + padding: 9px 14px; + background: var(--bg-secondary); + border: 1px solid var(--border-color); + border-radius: 10px; + cursor: pointer; + user-select: none; + transition: background 0.15s; +} +.ac-phase.open .ac-phase-header { + border-radius: 10px 10px 0 0; +} +.ac-phase-header:hover { + background: var(--bg-tertiary); +} + +.ac-phase-num { + font-family: var(--font-mono, monospace); + font-size: 10px; + font-weight: 600; + color: var(--accent, #3b82f6); + background: rgba(59,130,246,0.08); + padding: 2px 8px; + border-radius: 4px; + letter-spacing: 0.04em; + white-space: nowrap; + border: 1px solid rgba(59,130,246,0.1); +} + +.ac-phase-title { + font-family: var(--font-display); + font-size: 13px; + font-weight: 600; + color: var(--text-primary); + flex: 1; +} + +.ac-phase-dots { + display: flex; + gap: 3px; + align-items: center; +} +.ac-phase-dot { + width: 6px; + height: 6px; + border-radius: 50%; + flex-shrink: 0; +} +.ac-phase-dot.completed { background: var(--success, #16a34a); } +.ac-phase-dot.failed { background: var(--danger, #dc2626); } +.ac-phase-dot.running { background: var(--warning, #d97706); animation: ac-pulse 1.5s ease-in-out infinite; } +.ac-phase-dot.pending { background: var(--text-tertiary, #6b7280); opacity: 0.4; } + +@keyframes ac-pulse { + 0%, 100% { opacity: 1; } + 50% { opacity: 0.35; } +} + +.ac-phase-meta { + display: flex; + align-items: center; + gap: 12px; + font-family: var(--font-mono, monospace); + font-size: 11px; + color: var(--text-secondary); +} +.ac-phase-meta .findings-ct { color: var(--danger, #dc2626); font-weight: 600; } +.ac-phase-meta .running-ct { color: var(--warning, #d97706); font-weight: 500; } + +.ac-phase-chevron { + color: var(--text-tertiary, #6b7280); + font-size: 11px; + transition: transform 0.25s cubic-bezier(0.16,1,0.3,1); + width: 14px; + text-align: center; +} +.ac-phase.open .ac-phase-chevron { + transform: rotate(90deg); +} + +.ac-phase-body { + max-height: 0; + overflow: hidden; + transition: max-height 0.35s cubic-bezier(0.16,1,0.3,1); + background: var(--bg-secondary); + border-left: 1px solid var(--border-color); + border-right: 1px solid var(--border-color); + border-bottom: 1px solid var(--border-color); + border-radius: 0 0 10px 10px; +} +.ac-phase.open .ac-phase-body { + max-height: 2000px; +} +.ac-phase-body-inner { + padding: 4px 6px; + display: flex; + flex-direction: column; + gap: 1px; +} + +/* Tool rows */ +.ac-tool-row { + display: grid; + grid-template-columns: 5px 26px 1fr auto auto auto; + align-items: center; + gap: 8px; + padding: 7px 10px; + border-radius: 6px; + cursor: pointer; + transition: background 0.12s; +} +.ac-tool-row:hover { + background: rgba(255,255,255,0.02); +} +.ac-tool-row.expanded { + background: rgba(59,130,246,0.03); +} +.ac-tool-row.is-pending { + opacity: 0.45; + cursor: default; +} + +.ac-status-bar { + width: 4px; + height: 26px; + border-radius: 2px; + flex-shrink: 0; +} +.ac-status-bar.completed { background: var(--success, #16a34a); } +.ac-status-bar.failed { background: var(--danger, #dc2626); } +.ac-status-bar.running { background: var(--warning, #d97706); animation: ac-pulse 1.5s ease-in-out infinite; } +.ac-status-bar.pending { background: var(--text-tertiary, #6b7280); opacity: 0.25; } + +.ac-tool-icon { + font-size: 17px; + text-align: center; + line-height: 1; +} +.ac-tool-info { min-width: 0; } +.ac-tool-name { + font-size: 12.5px; + font-weight: 600; + color: var(--text-primary); + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; +} + +/* Category chips */ +.ac-cat-chip { + font-family: var(--font-mono, monospace); + font-size: 9px; + font-weight: 500; + padding: 1px 6px; + border-radius: 3px; + display: inline-block; + letter-spacing: 0.02em; +} +.ac-cat-chip.recon { color: #38bdf8; background: rgba(56,189,248,0.1); } +.ac-cat-chip.api { color: #818cf8; background: rgba(129,140,248,0.1); } +.ac-cat-chip.headers { color: #06b6d4; background: rgba(6,182,212,0.1); } +.ac-cat-chip.csp { color: #d946ef; background: rgba(217,70,239,0.1); } +.ac-cat-chip.cookies { color: #f59e0b; background: rgba(245,158,11,0.1); } +.ac-cat-chip.logs { color: #78716c; background: rgba(120,113,108,0.1); } +.ac-cat-chip.ratelimit { color: #64748b; background: rgba(100,116,139,0.1); } +.ac-cat-chip.cors { color: #8b5cf6; background: rgba(139,92,246,0.1); } +.ac-cat-chip.tls { color: #14b8a6; background: rgba(20,184,166,0.1); } +.ac-cat-chip.redirect { color: #fb923c; background: rgba(251,146,60,0.1); } +.ac-cat-chip.email { color: #0ea5e9; background: rgba(14,165,233,0.1); } +.ac-cat-chip.auth { color: #f43f5e; background: rgba(244,63,94,0.1); } +.ac-cat-chip.xss { color: #f97316; background: rgba(249,115,22,0.1); } +.ac-cat-chip.sqli { color: #ef4444; background: rgba(239,68,68,0.1); } +.ac-cat-chip.ssrf { color: #a855f7; background: rgba(168,85,247,0.1); } +.ac-cat-chip.idor { color: #ec4899; background: rgba(236,72,153,0.1); } +.ac-cat-chip.fuzzer { color: #a78bfa; background: rgba(167,139,250,0.1); } +.ac-cat-chip.cve { color: #dc2626; background: rgba(220,38,38,0.1); } +.ac-cat-chip.default { color: #94a3b8; background: rgba(148,163,184,0.1); } + +.ac-tool-duration { + font-family: var(--font-mono, monospace); + font-size: 10px; + color: var(--text-tertiary, #6b7280); + white-space: nowrap; + min-width: 48px; + text-align: right; +} +.ac-tool-duration.running-text { + color: var(--warning, #d97706); + font-weight: 500; +} + +.ac-findings-pill { + font-family: var(--font-mono, monospace); + font-size: 10px; + font-weight: 700; + display: inline-flex; + align-items: center; + justify-content: center; + min-width: 22px; + padding: 1px 7px; + border-radius: 9px; + line-height: 1.4; + text-align: center; +} +.ac-findings-pill.has { background: rgba(220,38,38,0.12); color: var(--danger, #dc2626); } +.ac-findings-pill.zero { background: transparent; color: var(--text-tertiary, #6b7280); font-weight: 400; opacity: 0.5; } + +.ac-risk-val { + font-family: var(--font-mono, monospace); + font-size: 10px; + font-weight: 700; + min-width: 32px; + text-align: right; +} +.ac-risk-val.high { color: var(--danger, #dc2626); } +.ac-risk-val.medium { color: var(--warning, #d97706); } +.ac-risk-val.low { color: var(--text-secondary); } +.ac-risk-val.none { color: transparent; } + +/* Tool detail (expanded) */ +.ac-tool-detail { + max-height: 0; + overflow: hidden; + transition: max-height 0.28s cubic-bezier(0.16,1,0.3,1); +} +.ac-tool-detail.open { + max-height: 300px; +} +.ac-tool-detail-inner { + padding: 6px 10px 10px 49px; + font-size: 12px; + line-height: 1.55; + color: var(--text-secondary); +} +.ac-reasoning-block { + background: rgba(59,130,246,0.03); + border-left: 2px solid var(--accent, #3b82f6); + padding: 7px 12px; + border-radius: 0 6px 6px 0; + font-style: italic; + margin-bottom: 8px; + color: var(--text-secondary); +} +.ac-detail-grid { + display: grid; + grid-template-columns: auto 1fr; + gap: 3px 14px; + font-family: var(--font-mono, monospace); + font-size: 10px; +} +.ac-detail-label { + color: var(--text-tertiary, #6b7280); + text-transform: uppercase; + font-size: 9px; + letter-spacing: 0.04em; +} +.ac-detail-value { + color: var(--text-secondary); +} diff --git a/compliance-dashboard/src/app.rs b/compliance-dashboard/src/app.rs index 91e382f..eb850ea 100644 --- a/compliance-dashboard/src/app.rs +++ b/compliance-dashboard/src/app.rs @@ -38,6 +38,10 @@ pub enum Route { DastFindingsPage {}, #[route("/dast/findings/:id")] DastFindingDetailPage { id: String }, + #[route("/pentest")] + PentestDashboardPage {}, + #[route("/pentest/:session_id")] + PentestSessionPage { session_id: String }, #[route("/mcp-servers")] McpServersPage {}, #[route("/settings")] @@ -49,7 +53,6 @@ const MAIN_CSS: Asset = asset!("/assets/main.css"); const TAILWIND_CSS: Asset = asset!("/assets/tailwind.css"); const VIS_NETWORK_JS: Asset = asset!("/assets/vis-network.min.js"); const GRAPH_VIZ_JS: Asset = asset!("/assets/graph-viz.js"); - #[component] pub fn App() -> Element { rsx! { diff --git a/compliance-dashboard/src/components/sidebar.rs b/compliance-dashboard/src/components/sidebar.rs index 4356c1a..4522fae 100644 --- a/compliance-dashboard/src/components/sidebar.rs +++ b/compliance-dashboard/src/components/sidebar.rs @@ -47,6 +47,11 @@ pub fn Sidebar() -> Element { route: Route::DastOverviewPage {}, icon: rsx! { Icon { icon: BsBug, width: 18, height: 18 } }, }, + NavItem { + label: "Pentest", + route: Route::PentestDashboardPage {}, + icon: rsx! { Icon { icon: BsLightningCharge, width: 18, height: 18 } }, + }, NavItem { label: "Settings", route: Route::SettingsPage {}, @@ -78,6 +83,7 @@ pub fn Sidebar() -> Element { (Route::DastTargetsPage {}, Route::DastOverviewPage {}) => true, (Route::DastFindingsPage {}, Route::DastOverviewPage {}) => true, (Route::DastFindingDetailPage { .. }, Route::DastOverviewPage {}) => true, + (Route::PentestSessionPage { .. }, Route::PentestDashboardPage {}) => true, (a, b) => a == b, }; let class = if is_active { "nav-item active" } else { "nav-item" }; diff --git a/compliance-dashboard/src/infrastructure/mod.rs b/compliance-dashboard/src/infrastructure/mod.rs index 1033ae9..490c63b 100644 --- a/compliance-dashboard/src/infrastructure/mod.rs +++ b/compliance-dashboard/src/infrastructure/mod.rs @@ -7,6 +7,7 @@ pub mod findings; pub mod graph; pub mod issues; pub mod mcp; +pub mod pentest; #[allow(clippy::too_many_arguments)] pub mod repositories; pub mod sbom; diff --git a/compliance-dashboard/src/infrastructure/pentest.rs b/compliance-dashboard/src/infrastructure/pentest.rs new file mode 100644 index 0000000..a9605f2 --- /dev/null +++ b/compliance-dashboard/src/infrastructure/pentest.rs @@ -0,0 +1,308 @@ +use dioxus::prelude::*; +use serde::{Deserialize, Serialize}; + +use super::dast::DastFindingsResponse; + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct PentestSessionsResponse { + pub data: Vec, + pub total: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct PentestSessionResponse { + pub data: serde_json::Value, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct PentestMessagesResponse { + pub data: Vec, + pub total: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct PentestStatsResponse { + pub data: serde_json::Value, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct AttackChainResponse { + pub data: Vec, +} + +#[server] +pub async fn fetch_pentest_sessions() -> Result { + let state: super::server_state::ServerState = + dioxus_fullstack::FullstackContext::extract().await?; + + // Fetch sessions + let url = format!("{}/api/v1/pentest/sessions", state.agent_api_url); + let resp = reqwest::get(&url) + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + let mut body: PentestSessionsResponse = resp + .json() + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + + // Fetch DAST targets to resolve target names + let targets_url = format!("{}/api/v1/dast/targets", state.agent_api_url); + if let Ok(tresp) = reqwest::get(&targets_url).await { + if let Ok(tbody) = tresp.json::().await { + let targets = tbody.get("data").and_then(|v| v.as_array()); + if let Some(targets) = targets { + // Build target_id -> name lookup + let target_map: std::collections::HashMap = targets + .iter() + .filter_map(|t| { + let id = t.get("_id")?.get("$oid")?.as_str()?.to_string(); + let name = t.get("name")?.as_str()?.to_string(); + Some((id, name)) + }) + .collect(); + + // Enrich sessions with target_name + for session in body.data.iter_mut() { + if let Some(tid) = session.get("target_id").and_then(|v| v.as_str()) { + if let Some(name) = target_map.get(tid) { + session.as_object_mut().map(|obj| { + obj.insert( + "target_name".to_string(), + serde_json::Value::String(name.clone()), + ) + }); + } + } + } + } + } + } + + Ok(body) +} + +#[server] +pub async fn fetch_pentest_session(id: String) -> Result { + let state: super::server_state::ServerState = + dioxus_fullstack::FullstackContext::extract().await?; + let url = format!("{}/api/v1/pentest/sessions/{id}", state.agent_api_url); + let resp = reqwest::get(&url) + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + let mut body: PentestSessionResponse = resp + .json() + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + + // Resolve target name from targets list + if let Some(tid) = body.data.get("target_id").and_then(|v| v.as_str()) { + let targets_url = format!("{}/api/v1/dast/targets", state.agent_api_url); + if let Ok(tresp) = reqwest::get(&targets_url).await { + if let Ok(tbody) = tresp.json::().await { + if let Some(targets) = tbody.get("data").and_then(|v| v.as_array()) { + for t in targets { + let t_id = t.get("_id").and_then(|v| v.get("$oid")).and_then(|v| v.as_str()).unwrap_or(""); + if t_id == tid { + if let Some(name) = t.get("name").and_then(|v| v.as_str()) { + body.data.as_object_mut().map(|obj| { + obj.insert("target_name".to_string(), serde_json::Value::String(name.to_string())) + }); + } + break; + } + } + } + } + } + } + + Ok(body) +} + +#[server] +pub async fn fetch_pentest_messages( + session_id: String, +) -> Result { + let state: super::server_state::ServerState = + dioxus_fullstack::FullstackContext::extract().await?; + let url = format!( + "{}/api/v1/pentest/sessions/{session_id}/messages", + state.agent_api_url + ); + let resp = reqwest::get(&url) + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + let body: PentestMessagesResponse = resp + .json() + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + Ok(body) +} + +#[server] +pub async fn fetch_pentest_stats() -> Result { + let state: super::server_state::ServerState = + dioxus_fullstack::FullstackContext::extract().await?; + let url = format!("{}/api/v1/pentest/stats", state.agent_api_url); + let resp = reqwest::get(&url) + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + let body: PentestStatsResponse = resp + .json() + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + Ok(body) +} + +#[server] +pub async fn fetch_attack_chain( + session_id: String, +) -> Result { + let state: super::server_state::ServerState = + dioxus_fullstack::FullstackContext::extract().await?; + let url = format!( + "{}/api/v1/pentest/sessions/{session_id}/attack-chain", + state.agent_api_url + ); + let resp = reqwest::get(&url) + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + let body: AttackChainResponse = resp + .json() + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + Ok(body) +} + +#[server] +pub async fn create_pentest_session( + target_id: String, + strategy: String, + message: String, +) -> Result { + let state: super::server_state::ServerState = + dioxus_fullstack::FullstackContext::extract().await?; + let url = format!("{}/api/v1/pentest/sessions", state.agent_api_url); + let client = reqwest::Client::new(); + let resp = client + .post(&url) + .json(&serde_json::json!({ + "target_id": target_id, + "strategy": strategy, + "message": message, + })) + .send() + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + let body: PentestSessionResponse = resp + .json() + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + Ok(body) +} + +#[server] +pub async fn send_pentest_message( + session_id: String, + message: String, +) -> Result { + let state: super::server_state::ServerState = + dioxus_fullstack::FullstackContext::extract().await?; + let url = format!( + "{}/api/v1/pentest/sessions/{session_id}/chat", + state.agent_api_url + ); + let client = reqwest::Client::new(); + let resp = client + .post(&url) + .json(&serde_json::json!({ + "message": message, + })) + .send() + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + let body: PentestMessagesResponse = resp + .json() + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + Ok(body) +} + +#[server] +pub async fn stop_pentest_session(session_id: String) -> Result<(), ServerFnError> { + let state: super::server_state::ServerState = + dioxus_fullstack::FullstackContext::extract().await?; + let url = format!( + "{}/api/v1/pentest/sessions/{session_id}/stop", + state.agent_api_url + ); + let client = reqwest::Client::new(); + client + .post(&url) + .send() + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + Ok(()) +} + +#[server] +pub async fn fetch_pentest_findings( + session_id: String, +) -> Result { + let state: super::server_state::ServerState = + dioxus_fullstack::FullstackContext::extract().await?; + let url = format!( + "{}/api/v1/pentest/sessions/{session_id}/findings", + state.agent_api_url + ); + let resp = reqwest::get(&url) + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + let body: DastFindingsResponse = resp + .json() + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + Ok(body) +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct ExportReportResponse { + pub archive_base64: String, + pub sha256: String, + pub filename: String, +} + +#[server] +pub async fn export_pentest_report( + session_id: String, + password: String, + requester_name: String, + requester_email: String, +) -> Result { + let state: super::server_state::ServerState = + dioxus_fullstack::FullstackContext::extract().await?; + let url = format!( + "{}/api/v1/pentest/sessions/{session_id}/export", + state.agent_api_url + ); + let client = reqwest::Client::new(); + let resp = client + .post(&url) + .json(&serde_json::json!({ + "password": password, + "requester_name": requester_name, + "requester_email": requester_email, + })) + .send() + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + if !resp.status().is_success() { + let text = resp.text().await.unwrap_or_default(); + return Err(ServerFnError::new(format!("Export failed: {text}"))); + } + let body: ExportReportResponse = resp + .json() + .await + .map_err(|e| ServerFnError::new(e.to_string()))?; + Ok(body) +} diff --git a/compliance-dashboard/src/pages/dast_findings.rs b/compliance-dashboard/src/pages/dast_findings.rs index 4fa4cac..422f183 100644 --- a/compliance-dashboard/src/pages/dast_findings.rs +++ b/compliance-dashboard/src/pages/dast_findings.rs @@ -11,6 +11,11 @@ use crate::infrastructure::dast::fetch_dast_findings; pub fn DastFindingsPage() -> Element { let findings = use_resource(|| async { fetch_dast_findings().await.ok() }); + let mut filter_severity = use_signal(|| "all".to_string()); + let mut filter_vuln_type = use_signal(|| "all".to_string()); + let mut filter_exploitable = use_signal(|| "all".to_string()); + let mut search_text = use_signal(String::new); + rsx! { div { class: "back-nav", button { @@ -26,14 +31,105 @@ pub fn DastFindingsPage() -> Element { description: "Vulnerabilities discovered through dynamic application security testing", } + // Filter bar + div { style: "display: flex; gap: 10px; margin-bottom: 12px; flex-wrap: wrap; align-items: center;", + // Search + div { style: "flex: 1; min-width: 180px;", + input { + class: "chat-input", + style: "width: 100%; padding: 6px 10px; font-size: 0.85rem;", + placeholder: "Search title or endpoint...", + value: "{search_text}", + oninput: move |e| search_text.set(e.value()), + } + } + // Severity + select { + style: "padding: 6px 10px; border-radius: 6px; border: 1px solid var(--border-color); background: var(--bg-secondary); color: var(--text-primary); font-size: 0.85rem;", + value: "{filter_severity}", + onchange: move |e| filter_severity.set(e.value()), + option { value: "all", "All Severities" } + option { value: "critical", "Critical" } + option { value: "high", "High" } + option { value: "medium", "Medium" } + option { value: "low", "Low" } + option { value: "info", "Info" } + } + // Vuln type + select { + style: "padding: 6px 10px; border-radius: 6px; border: 1px solid var(--border-color); background: var(--bg-secondary); color: var(--text-primary); font-size: 0.85rem;", + value: "{filter_vuln_type}", + onchange: move |e| filter_vuln_type.set(e.value()), + option { value: "all", "All Types" } + option { value: "sql_injection", "SQL Injection" } + option { value: "xss", "XSS" } + option { value: "auth_bypass", "Auth Bypass" } + option { value: "ssrf", "SSRF" } + option { value: "api_misconfiguration", "API Misconfiguration" } + option { value: "open_redirect", "Open Redirect" } + option { value: "idor", "IDOR" } + option { value: "information_disclosure", "Information Disclosure" } + option { value: "security_misconfiguration", "Security Misconfiguration" } + option { value: "broken_auth", "Broken Auth" } + option { value: "dns_misconfiguration", "DNS Misconfiguration" } + option { value: "email_security", "Email Security" } + option { value: "tls_misconfiguration", "TLS Misconfiguration" } + option { value: "cookie_security", "Cookie Security" } + option { value: "csp_issue", "CSP Issue" } + option { value: "cors_misconfiguration", "CORS Misconfiguration" } + option { value: "rate_limit_absent", "Rate Limit Absent" } + option { value: "console_log_leakage", "Console Log Leakage" } + option { value: "security_header_missing", "Security Header Missing" } + option { value: "known_cve_exploit", "Known CVE Exploit" } + option { value: "other", "Other" } + } + // Exploitable + select { + style: "padding: 6px 10px; border-radius: 6px; border: 1px solid var(--border-color); background: var(--bg-secondary); color: var(--text-primary); font-size: 0.85rem;", + value: "{filter_exploitable}", + onchange: move |e| filter_exploitable.set(e.value()), + option { value: "all", "All" } + option { value: "yes", "Exploitable" } + option { value: "no", "Unconfirmed" } + } + } + div { class: "card", match &*findings.read() { Some(Some(data)) => { - let finding_list = &data.data; - if finding_list.is_empty() { - rsx! { p { "No DAST findings yet. Run a scan to discover vulnerabilities." } } + let sev_filter = filter_severity.read().clone(); + let vt_filter = filter_vuln_type.read().clone(); + let exp_filter = filter_exploitable.read().clone(); + let search = search_text.read().to_lowercase(); + + let filtered: Vec<_> = data.data.iter().filter(|f| { + let severity = f.get("severity").and_then(|v| v.as_str()).unwrap_or("info"); + let vuln_type = f.get("vuln_type").and_then(|v| v.as_str()).unwrap_or(""); + let exploitable = f.get("exploitable").and_then(|v| v.as_bool()).unwrap_or(false); + let title = f.get("title").and_then(|v| v.as_str()).unwrap_or("").to_lowercase(); + let endpoint = f.get("endpoint").and_then(|v| v.as_str()).unwrap_or("").to_lowercase(); + + (sev_filter == "all" || severity == sev_filter) + && (vt_filter == "all" || vuln_type == vt_filter) + && match exp_filter.as_str() { + "yes" => exploitable, + "no" => !exploitable, + _ => true, + } + && (search.is_empty() || title.contains(&search) || endpoint.contains(&search)) + }).collect(); + + if filtered.is_empty() { + if data.data.is_empty() { + rsx! { p { style: "padding: 16px;", "No DAST findings yet. Run a scan to discover vulnerabilities." } } + } else { + rsx! { p { style: "padding: 16px; color: var(--text-secondary);", "No findings match the current filters." } } + } } else { rsx! { + div { style: "padding: 8px 16px; font-size: 0.8rem; color: var(--text-secondary);", + "Showing {filtered.len()} of {data.data.len()} findings" + } table { class: "table", thead { tr { @@ -46,7 +142,7 @@ pub fn DastFindingsPage() -> Element { } } tbody { - for finding in finding_list { + for finding in filtered { { let id = finding.get("_id").and_then(|v| v.get("$oid")).and_then(|v| v.as_str()).unwrap_or("").to_string(); let severity = finding.get("severity").and_then(|v| v.as_str()).unwrap_or("info").to_string(); diff --git a/compliance-dashboard/src/pages/mod.rs b/compliance-dashboard/src/pages/mod.rs index 623ec4a..bdc9281 100644 --- a/compliance-dashboard/src/pages/mod.rs +++ b/compliance-dashboard/src/pages/mod.rs @@ -12,6 +12,8 @@ pub mod impact_analysis; pub mod issues; pub mod mcp_servers; pub mod overview; +pub mod pentest_dashboard; +pub mod pentest_session; pub mod repositories; pub mod sbom; pub mod settings; @@ -30,6 +32,8 @@ pub use impact_analysis::ImpactAnalysisPage; pub use issues::IssuesPage; pub use mcp_servers::McpServersPage; pub use overview::OverviewPage; +pub use pentest_dashboard::PentestDashboardPage; +pub use pentest_session::PentestSessionPage; pub use repositories::RepositoriesPage; pub use sbom::SbomPage; pub use settings::SettingsPage; diff --git a/compliance-dashboard/src/pages/pentest_dashboard.rs b/compliance-dashboard/src/pages/pentest_dashboard.rs new file mode 100644 index 0000000..a8643b4 --- /dev/null +++ b/compliance-dashboard/src/pages/pentest_dashboard.rs @@ -0,0 +1,398 @@ +use dioxus::prelude::*; +use dioxus_free_icons::icons::bs_icons::*; +use dioxus_free_icons::Icon; + +use crate::app::Route; +use crate::components::page_header::PageHeader; +use crate::infrastructure::dast::fetch_dast_targets; +use crate::infrastructure::pentest::{ + create_pentest_session, fetch_pentest_sessions, fetch_pentest_stats, stop_pentest_session, +}; + +#[component] +pub fn PentestDashboardPage() -> Element { + let mut sessions = use_resource(|| async { fetch_pentest_sessions().await.ok() }); + let stats = use_resource(|| async { fetch_pentest_stats().await.ok() }); + let targets = use_resource(|| async { fetch_dast_targets().await.ok() }); + + let mut show_modal = use_signal(|| false); + let mut new_target_id = use_signal(String::new); + let mut new_strategy = use_signal(|| "comprehensive".to_string()); + let mut new_message = use_signal(String::new); + let mut creating = use_signal(|| false); + + let on_create = move |_| { + let tid = new_target_id.read().clone(); + let strat = new_strategy.read().clone(); + let msg = new_message.read().clone(); + if tid.is_empty() || msg.is_empty() { + return; + } + creating.set(true); + spawn(async move { + match create_pentest_session(tid, strat, msg).await { + Ok(resp) => { + let session_id = resp + .data + .get("_id") + .and_then(|v| v.get("$oid")) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + creating.set(false); + show_modal.set(false); + new_target_id.set(String::new()); + new_message.set(String::new()); + if !session_id.is_empty() { + navigator().push(Route::PentestSessionPage { + session_id: session_id.clone(), + }); + } else { + sessions.restart(); + } + } + Err(_) => { + creating.set(false); + } + } + }); + }; + + // Extract stats values + let running_sessions = { + let s = stats.read(); + match &*s { + Some(Some(data)) => data + .data + .get("running_sessions") + .and_then(|v| v.as_u64()) + .unwrap_or(0), + _ => 0, + } + }; + let total_vulns = { + let s = stats.read(); + match &*s { + Some(Some(data)) => data + .data + .get("total_vulnerabilities") + .and_then(|v| v.as_u64()) + .unwrap_or(0), + _ => 0, + } + }; + let tool_invocations = { + let s = stats.read(); + match &*s { + Some(Some(data)) => data + .data + .get("total_tool_invocations") + .and_then(|v| v.as_u64()) + .unwrap_or(0), + _ => 0, + } + }; + let success_rate = { + let s = stats.read(); + match &*s { + Some(Some(data)) => data + .data + .get("tool_success_rate") + .and_then(|v| v.as_f64()) + .unwrap_or(0.0), + _ => 0.0, + } + }; + + // Severity counts from stats (nested under severity_distribution) + let sev_dist = { + let s = stats.read(); + match &*s { + Some(Some(data)) => data + .data + .get("severity_distribution") + .cloned() + .unwrap_or(serde_json::Value::Null), + _ => serde_json::Value::Null, + } + }; + let severity_critical = sev_dist.get("critical").and_then(|v| v.as_u64()).unwrap_or(0); + let severity_high = sev_dist.get("high").and_then(|v| v.as_u64()).unwrap_or(0); + let severity_medium = sev_dist.get("medium").and_then(|v| v.as_u64()).unwrap_or(0); + let severity_low = sev_dist.get("low").and_then(|v| v.as_u64()).unwrap_or(0); + + rsx! { + PageHeader { + title: "Pentest Dashboard", + description: "AI-powered penetration testing sessions — autonomous security assessment", + } + + // Stat cards + div { class: "stat-cards", style: "margin-bottom: 24px;", + div { class: "stat-card-item", + div { class: "stat-card-value", "{running_sessions}" } + div { class: "stat-card-label", + Icon { icon: BsPlayCircle, width: 14, height: 14 } + " Running Sessions" + } + } + div { class: "stat-card-item", + div { class: "stat-card-value", "{total_vulns}" } + div { class: "stat-card-label", + Icon { icon: BsShieldExclamation, width: 14, height: 14 } + " Total Vulnerabilities" + } + } + div { class: "stat-card-item", + div { class: "stat-card-value", "{tool_invocations}" } + div { class: "stat-card-label", + Icon { icon: BsWrench, width: 14, height: 14 } + " Tool Invocations" + } + } + div { class: "stat-card-item", + div { class: "stat-card-value", "{success_rate:.0}%" } + div { class: "stat-card-label", + Icon { icon: BsCheckCircle, width: 14, height: 14 } + " Success Rate" + } + } + } + + // Severity distribution + div { class: "card", style: "margin-bottom: 24px; padding: 16px;", + div { style: "display: flex; align-items: center; gap: 16px; flex-wrap: wrap;", + span { style: "font-weight: 600; color: var(--text-secondary); font-size: 0.85rem;", "Severity Distribution" } + span { + class: "badge", + style: "background: #dc2626; color: #fff;", + "Critical: {severity_critical}" + } + span { + class: "badge", + style: "background: #ea580c; color: #fff;", + "High: {severity_high}" + } + span { + class: "badge", + style: "background: #d97706; color: #fff;", + "Medium: {severity_medium}" + } + span { + class: "badge", + style: "background: #2563eb; color: #fff;", + "Low: {severity_low}" + } + } + } + + // Actions row + div { style: "display: flex; gap: 12px; margin-bottom: 24px;", + button { + class: "btn btn-primary", + onclick: move |_| show_modal.set(true), + Icon { icon: BsPlusCircle, width: 14, height: 14 } + " New Pentest" + } + } + + // Sessions list + div { class: "card", + div { class: "card-header", "Recent Pentest Sessions" } + match &*sessions.read() { + Some(Some(data)) => { + let sess_list = &data.data; + if sess_list.is_empty() { + rsx! { + div { style: "padding: 32px; text-align: center; color: var(--text-secondary);", + p { "No pentest sessions yet. Start one to begin autonomous security testing." } + } + } + } else { + rsx! { + div { style: "display: grid; gap: 12px; padding: 16px;", + for session in sess_list { + { + let id = session.get("_id") + .and_then(|v| v.get("$oid")) + .and_then(|v| v.as_str()) + .unwrap_or("-").to_string(); + let target_name = session.get("target_name").and_then(|v| v.as_str()).unwrap_or("Unknown Target").to_string(); + let status = session.get("status").and_then(|v| v.as_str()).unwrap_or("unknown").to_string(); + let strategy = session.get("strategy").and_then(|v| v.as_str()).unwrap_or("-").to_string(); + let findings_count = session.get("findings_count").and_then(|v| v.as_u64()).unwrap_or(0); + let tool_count = session.get("tool_invocations").and_then(|v| v.as_u64()).unwrap_or(0); + let created_at = session.get("created_at").and_then(|v| v.as_str()).unwrap_or("-").to_string(); + let status_style = match status.as_str() { + "running" => "background: #16a34a; color: #fff;", + "completed" => "background: #2563eb; color: #fff;", + "failed" => "background: #dc2626; color: #fff;", + "paused" => "background: #d97706; color: #fff;", + _ => "background: var(--bg-tertiary); color: var(--text-secondary);", + }; + { + let is_session_running = status == "running"; + let stop_id = id.clone(); + rsx! { + div { class: "card", style: "padding: 16px; transition: border-color 0.15s;", + Link { + to: Route::PentestSessionPage { session_id: id.clone() }, + style: "text-decoration: none; cursor: pointer; display: block;", + div { style: "display: flex; justify-content: space-between; align-items: flex-start;", + div { + div { style: "font-weight: 600; font-size: 1rem; margin-bottom: 4px; color: var(--text-primary);", + "{target_name}" + } + div { style: "display: flex; gap: 8px; align-items: center; flex-wrap: wrap;", + span { + class: "badge", + style: "{status_style}", + "{status}" + } + span { + class: "badge", + style: "background: var(--bg-tertiary); color: var(--text-secondary);", + "{strategy}" + } + } + } + div { style: "text-align: right; font-size: 0.85rem; color: var(--text-secondary);", + div { style: "margin-bottom: 4px;", + Icon { icon: BsShieldExclamation, width: 12, height: 12 } + " {findings_count} findings" + } + div { style: "margin-bottom: 4px;", + Icon { icon: BsWrench, width: 12, height: 12 } + " {tool_count} tools" + } + div { "{created_at}" } + } + } + } + if is_session_running { + div { style: "margin-top: 8px; display: flex; justify-content: flex-end;", + button { + class: "btn btn-ghost", + style: "font-size: 0.8rem; padding: 4px 12px; color: #dc2626; border-color: #dc2626;", + onclick: move |e| { + e.stop_propagation(); + e.prevent_default(); + let sid = stop_id.clone(); + spawn(async move { + let _ = stop_pentest_session(sid).await; + sessions.restart(); + }); + }, + Icon { icon: BsStopCircle, width: 12, height: 12 } + " Stop" + } + } + } + } + } + } + } + } + } + } + } + }, + Some(None) => rsx! { p { style: "padding: 16px;", "Failed to load sessions." } }, + None => rsx! { p { style: "padding: 16px;", "Loading..." } }, + } + } + + // New Pentest Modal + if *show_modal.read() { + div { + style: "position: fixed; inset: 0; background: rgba(0,0,0,0.6); display: flex; align-items: center; justify-content: center; z-index: 1000;", + onclick: move |_| show_modal.set(false), + div { + style: "background: var(--bg-secondary); border: 1px solid var(--border-color); border-radius: 12px; padding: 24px; width: 480px; max-width: 90vw;", + onclick: move |e| e.stop_propagation(), + h3 { style: "margin: 0 0 16px 0;", "New Pentest Session" } + + // Target selection + div { style: "margin-bottom: 12px;", + label { style: "display: block; font-size: 0.85rem; color: var(--text-secondary); margin-bottom: 4px;", + "Target" + } + select { + class: "chat-input", + style: "width: 100%; padding: 8px; resize: none; height: auto;", + value: "{new_target_id}", + onchange: move |e| new_target_id.set(e.value()), + option { value: "", "Select a target..." } + match &*targets.read() { + Some(Some(data)) => { + rsx! { + for target in &data.data { + { + let tid = target.get("_id") + .and_then(|v| v.get("$oid")) + .and_then(|v| v.as_str()) + .unwrap_or("").to_string(); + let tname = target.get("name").and_then(|v| v.as_str()).unwrap_or("Unknown").to_string(); + let turl = target.get("base_url").and_then(|v| v.as_str()).unwrap_or("").to_string(); + rsx! { + option { value: "{tid}", "{tname} ({turl})" } + } + } + } + } + }, + _ => rsx! {}, + } + } + } + + // Strategy selection + div { style: "margin-bottom: 12px;", + label { style: "display: block; font-size: 0.85rem; color: var(--text-secondary); margin-bottom: 4px;", + "Strategy" + } + select { + class: "chat-input", + style: "width: 100%; padding: 8px; resize: none; height: auto;", + value: "{new_strategy}", + onchange: move |e| new_strategy.set(e.value()), + option { value: "comprehensive", "Comprehensive" } + option { value: "quick", "Quick Scan" } + option { value: "owasp_top_10", "OWASP Top 10" } + option { value: "api_focused", "API Focused" } + option { value: "authentication", "Authentication" } + } + } + + // Initial message + div { style: "margin-bottom: 16px;", + label { style: "display: block; font-size: 0.85rem; color: var(--text-secondary); margin-bottom: 4px;", + "Initial Instructions" + } + textarea { + class: "chat-input", + style: "width: 100%; min-height: 80px;", + placeholder: "Describe the scope and goals of this pentest...", + value: "{new_message}", + oninput: move |e| new_message.set(e.value()), + } + } + + div { style: "display: flex; justify-content: flex-end; gap: 8px;", + button { + class: "btn btn-ghost", + onclick: move |_| show_modal.set(false), + "Cancel" + } + button { + class: "btn btn-primary", + disabled: *creating.read() || new_target_id.read().is_empty() || new_message.read().is_empty(), + onclick: on_create, + if *creating.read() { "Creating..." } else { "Start Pentest" } + } + } + } + } + } + } +} diff --git a/compliance-dashboard/src/pages/pentest_session.rs b/compliance-dashboard/src/pages/pentest_session.rs new file mode 100644 index 0000000..bc92112 --- /dev/null +++ b/compliance-dashboard/src/pages/pentest_session.rs @@ -0,0 +1,1141 @@ +use std::collections::{HashMap, VecDeque}; + +use dioxus::prelude::*; +use dioxus_free_icons::icons::bs_icons::*; +use dioxus_free_icons::Icon; + +use crate::app::Route; +use crate::components::severity_badge::SeverityBadge; +use crate::infrastructure::pentest::{ + export_pentest_report, fetch_attack_chain, fetch_pentest_findings, fetch_pentest_session, +}; + +#[component] +pub fn PentestSessionPage(session_id: String) -> Element { + let sid_for_session = session_id.clone(); + let sid_for_findings = session_id.clone(); + let sid_for_chain = session_id.clone(); + + let mut session = use_resource(move || { + let id = sid_for_session.clone(); + async move { fetch_pentest_session(id).await.ok() } + }); + let mut findings = use_resource(move || { + let id = sid_for_findings.clone(); + async move { fetch_pentest_findings(id).await.ok() } + }); + let mut attack_chain = use_resource(move || { + let id = sid_for_chain.clone(); + async move { fetch_attack_chain(id).await.ok() } + }); + + let mut active_tab = use_signal(|| "findings".to_string()); + let mut show_export_modal = use_signal(|| false); + let mut export_password = use_signal(String::new); + let mut exporting = use_signal(|| false); + let mut export_sha256 = use_signal(|| Option::::None); + let mut export_error = use_signal(|| Option::::None); + let mut poll_gen = use_signal(|| 0u32); + + // Extract session data + let session_data = session.read().clone(); + let sess = session_data.as_ref().and_then(|s| s.as_ref()); + + let session_status = sess + .and_then(|s| s.data.get("status")) + .and_then(|v| v.as_str()) + .unwrap_or("unknown") + .to_string(); + let target_name = sess + .and_then(|s| s.data.get("target_name")) + .and_then(|v| v.as_str()) + .unwrap_or("Pentest Session") + .to_string(); + let strategy = sess + .and_then(|s| s.data.get("strategy")) + .and_then(|v| v.as_str()) + .unwrap_or("-") + .to_string(); + let tool_invocations = sess + .and_then(|s| s.data.get("tool_invocations")) + .and_then(|v| v.as_u64()) + .unwrap_or(0); + let tool_successes = sess + .and_then(|s| s.data.get("tool_successes")) + .and_then(|v| v.as_u64()) + .unwrap_or(0); + let findings_count = { + let f = findings.read(); + match &*f { + Some(Some(data)) => data.total.unwrap_or(0), + _ => 0, + } + }; + let started_at = sess + .and_then(|s| s.data.get("started_at")) + .and_then(|v| v.as_str()) + .unwrap_or("-") + .to_string(); + let completed_at = sess + .and_then(|s| s.data.get("completed_at")) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let success_rate = if tool_invocations == 0 { + 100.0 + } else { + (tool_successes as f64 / tool_invocations as f64) * 100.0 + }; + + let is_running = session_status == "running"; + + // Poll while running + use_effect(move || { + let _gen = *poll_gen.read(); + if is_running { + spawn(async move { + #[cfg(feature = "web")] + gloo_timers::future::TimeoutFuture::new(3_000).await; + #[cfg(not(feature = "web"))] + tokio::time::sleep(std::time::Duration::from_secs(3)).await; + findings.restart(); + attack_chain.restart(); + session.restart(); + let next = poll_gen.peek().wrapping_add(1); + poll_gen.set(next); + }); + } + }); + + // Severity counts from findings data + let (sev_critical, sev_high, sev_medium, sev_low, sev_info, exploitable_count) = { + let f = findings.read(); + match &*f { + Some(Some(data)) => { + let list = &data.data; + let c = list + .iter() + .filter(|f| { + f.get("severity").and_then(|v| v.as_str()) == Some("critical") + }) + .count(); + let h = list + .iter() + .filter(|f| f.get("severity").and_then(|v| v.as_str()) == Some("high")) + .count(); + let m = list + .iter() + .filter(|f| { + f.get("severity").and_then(|v| v.as_str()) == Some("medium") + }) + .count(); + let l = list + .iter() + .filter(|f| f.get("severity").and_then(|v| v.as_str()) == Some("low")) + .count(); + let i = list + .iter() + .filter(|f| f.get("severity").and_then(|v| v.as_str()) == Some("info")) + .count(); + let e = list + .iter() + .filter(|f| { + f.get("exploitable").and_then(|v| v.as_bool()).unwrap_or(false) + }) + .count(); + (c, h, m, l, i, e) + } + _ => (0, 0, 0, 0, 0, 0), + } + }; + + let status_style = match session_status.as_str() { + "running" => "background: #16a34a; color: #fff;", + "completed" => "background: #2563eb; color: #fff;", + "failed" => "background: #dc2626; color: #fff;", + "paused" => "background: #d97706; color: #fff;", + _ => "background: var(--bg-tertiary); color: var(--text-secondary);", + }; + + // Export handler + let sid_for_export = session_id.clone(); + let do_export = move |_| { + let pw = export_password.read().clone(); + if pw.len() < 8 { + export_error.set(Some("Password must be at least 8 characters".to_string())); + return; + } + export_error.set(None); + export_sha256.set(None); + exporting.set(true); + let sid = sid_for_export.clone(); + spawn(async move { + // TODO: get real user info from auth context + match export_pentest_report( + sid.clone(), + pw, + String::new(), + String::new(), + ) + .await + { + Ok(resp) => { + export_sha256.set(Some(resp.sha256.clone())); + // Trigger download via JS + let js = format!( + r#" + try {{ + var raw = atob("{}"); + var bytes = new Uint8Array(raw.length); + for (var i = 0; i < raw.length; i++) bytes[i] = raw.charCodeAt(i); + var blob = new Blob([bytes], {{ type: "application/octet-stream" }}); + var url = URL.createObjectURL(blob); + var a = document.createElement("a"); + a.href = url; + a.download = "{}"; + document.body.appendChild(a); + a.click(); + document.body.removeChild(a); + URL.revokeObjectURL(url); + }} catch(e) {{ console.error("Download failed:", e); }} + "#, + resp.archive_base64, resp.filename, + ); + document::eval(&js); + } + Err(e) => { + export_error.set(Some(format!("{e}"))); + } + } + exporting.set(false); + }); + }; + + rsx! { + div { class: "back-nav", + Link { + to: Route::PentestDashboardPage {}, + class: "btn btn-ghost btn-back", + Icon { icon: BsArrowLeft, width: 16, height: 16 } + "Back to Pentest Dashboard" + } + } + + // Session header + div { style: "display: flex; align-items: center; justify-content: space-between; margin-bottom: 20px; flex-wrap: wrap; gap: 8px;", + div { + h2 { style: "margin: 0 0 4px 0;", "{target_name}" } + div { style: "display: flex; gap: 8px; align-items: center; flex-wrap: wrap;", + span { class: "badge", style: "{status_style}", "{session_status}" } + span { class: "badge", style: "background: var(--bg-tertiary); color: var(--text-secondary);", + "{strategy}" + } + if is_running { + span { style: "font-size: 0.8rem; color: var(--text-secondary);", + Icon { icon: BsPlayCircle, width: 12, height: 12 } + " Running..." + } + } + } + } + div { style: "display: flex; gap: 8px;", + button { + class: "btn btn-primary", + style: "font-size: 0.85rem;", + onclick: move |_| { + export_password.set(String::new()); + export_sha256.set(None); + export_error.set(None); + show_export_modal.set(true); + }, + Icon { icon: BsDownload, width: 14, height: 14 } + " Export Report" + } + } + } + + // Summary cards + div { class: "stat-cards", style: "margin-bottom: 20px;", + div { class: "stat-card-item", + div { class: "stat-card-value", "{findings_count}" } + div { class: "stat-card-label", + Icon { icon: BsShieldExclamation, width: 14, height: 14 } + " Findings" + } + } + div { class: "stat-card-item", + div { class: "stat-card-value", style: "color: #dc2626;", "{exploitable_count}" } + div { class: "stat-card-label", + Icon { icon: BsExclamationTriangle, width: 14, height: 14 } + " Exploitable" + } + } + div { class: "stat-card-item", + div { class: "stat-card-value", "{tool_invocations}" } + div { class: "stat-card-label", + Icon { icon: BsWrench, width: 14, height: 14 } + " Tool Invocations" + } + } + div { class: "stat-card-item", + div { class: "stat-card-value", "{success_rate:.0}%" } + div { class: "stat-card-label", + Icon { icon: BsCheckCircle, width: 14, height: 14 } + " Success Rate" + } + } + } + + // Severity distribution bar + div { class: "card", style: "margin-bottom: 20px; padding: 14px;", + div { style: "display: flex; align-items: center; gap: 14px; flex-wrap: wrap;", + span { style: "font-weight: 600; color: var(--text-secondary); font-size: 0.85rem;", "Severity Distribution" } + span { class: "badge", style: "background: #dc2626; color: #fff;", "Critical: {sev_critical}" } + span { class: "badge", style: "background: #ea580c; color: #fff;", "High: {sev_high}" } + span { class: "badge", style: "background: #d97706; color: #fff;", "Medium: {sev_medium}" } + span { class: "badge", style: "background: #2563eb; color: #fff;", "Low: {sev_low}" } + span { class: "badge", style: "background: #6b7280; color: #fff;", "Info: {sev_info}" } + } + } + + // Session details row + div { class: "card", style: "margin-bottom: 20px; padding: 14px;", + div { style: "display: grid; grid-template-columns: repeat(auto-fit, minmax(160px, 1fr)); gap: 12px; font-size: 0.85rem;", + div { + span { style: "color: var(--text-secondary);", "Started: " } + span { "{started_at}" } + } + if !completed_at.is_empty() { + div { + span { style: "color: var(--text-secondary);", "Completed: " } + span { "{completed_at}" } + } + } + div { + span { style: "color: var(--text-secondary);", "Tools: " } + span { "{tool_successes}/{tool_invocations} successful" } + } + } + } + + // Tabs: Findings / Attack Chain + div { class: "card", style: "overflow: hidden;", + div { style: "display: flex; border-bottom: 1px solid var(--border-color);", + button { + style: if *active_tab.read() == "findings" { + "flex: 1; padding: 12px; background: none; border: none; border-bottom: 2px solid #2563eb; color: var(--text-primary); cursor: pointer; font-weight: 600; font-size: 0.9rem;" + } else { + "flex: 1; padding: 12px; background: none; border: none; border-bottom: 2px solid transparent; color: var(--text-secondary); cursor: pointer; font-size: 0.9rem;" + }, + onclick: move |_| active_tab.set("findings".to_string()), + Icon { icon: BsShieldExclamation, width: 14, height: 14 } + " Findings ({findings_count})" + } + button { + style: if *active_tab.read() == "chain" { + "flex: 1; padding: 12px; background: none; border: none; border-bottom: 2px solid #2563eb; color: var(--text-primary); cursor: pointer; font-weight: 600; font-size: 0.9rem;" + } else { + "flex: 1; padding: 12px; background: none; border: none; border-bottom: 2px solid transparent; color: var(--text-secondary); cursor: pointer; font-size: 0.9rem;" + }, + onclick: move |_| { + active_tab.set("chain".to_string()); + }, + Icon { icon: BsDiagram3, width: 14, height: 14 } + " Attack Chain" + } + } + + // Tab content + div { style: "padding: 16px;", + if *active_tab.read() == "findings" { + // Findings list + match &*findings.read() { + Some(Some(data)) => { + let finding_list = &data.data; + if finding_list.is_empty() { + rsx! { + div { style: "text-align: center; color: var(--text-secondary); padding: 24px;", + if is_running { + p { "Scan in progress — findings will appear here." } + } else { + p { "No findings discovered." } + } + } + } + } else { + rsx! { + div { style: "display: flex; flex-direction: column; gap: 10px;", + for finding in finding_list { + { + let title = finding.get("title").and_then(|v| v.as_str()).unwrap_or("Untitled").to_string(); + let severity = finding.get("severity").and_then(|v| v.as_str()).unwrap_or("info").to_string(); + let vuln_type = finding.get("vuln_type").and_then(|v| v.as_str()).unwrap_or("-").to_string(); + let endpoint = finding.get("endpoint").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let method = finding.get("method").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let exploitable = finding.get("exploitable").and_then(|v| v.as_bool()).unwrap_or(false); + let description = finding.get("description").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let remediation = finding.get("remediation").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let cwe = finding.get("cwe").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let linked_sast = finding.get("linked_sast_finding_id").and_then(|v| v.as_str()).unwrap_or("").to_string(); + + rsx! { + div { style: "background: var(--bg-tertiary); border-radius: 8px; padding: 14px;", + // Header + div { style: "display: flex; justify-content: space-between; align-items: center; margin-bottom: 8px;", + div { style: "display: flex; align-items: center; gap: 8px;", + SeverityBadge { severity: severity } + span { style: "font-weight: 600; font-size: 0.95rem;", "{title}" } + } + div { style: "display: flex; gap: 4px;", + if exploitable { + span { class: "badge", style: "background: #dc2626; color: #fff; font-size: 0.7rem;", "Exploitable" } + } + span { class: "badge", style: "font-size: 0.7rem;", "{vuln_type}" } + } + } + // Endpoint + if !endpoint.is_empty() { + div { style: "font-family: monospace; font-size: 0.8rem; color: var(--text-secondary); margin-bottom: 6px;", + "{method} {endpoint}" + } + } + // CWE + if !cwe.is_empty() { + div { style: "font-size: 0.8rem; color: var(--text-secondary); margin-bottom: 4px;", + "CWE: {cwe}" + } + } + // Description + if !description.is_empty() { + div { style: "font-size: 0.85rem; margin-bottom: 8px; line-height: 1.5;", + "{description}" + } + } + // Remediation + if !remediation.is_empty() { + div { style: "font-size: 0.8rem; padding: 8px 10px; background: rgba(56, 189, 248, 0.08); border-left: 3px solid #38bdf8; border-radius: 0 4px 4px 0; margin-top: 6px;", + span { style: "font-weight: 600;", "Recommendation: " } + "{remediation}" + } + } + // Linked SAST + if !linked_sast.is_empty() { + div { style: "font-size: 0.75rem; color: var(--text-secondary); margin-top: 4px;", + "Correlated SAST finding: " + code { "{linked_sast}" } + } + } + } + } + } + } + } + } + } + }, + Some(None) => rsx! { p { style: "color: var(--text-secondary);", "Failed to load findings." } }, + None => rsx! { p { style: "color: var(--text-secondary);", "Loading..." } }, + } + } else { + // Attack chain visualization + match &*attack_chain.read() { + Some(Some(data)) => { + let steps = &data.data; + if steps.is_empty() { + rsx! { + div { style: "text-align: center; color: var(--text-secondary); padding: 24px;", + if is_running { + p { "Scan in progress — attack chain will appear here." } + } else { + p { "No attack chain steps recorded." } + } + } + } + } else { + rsx! { AttackChainView { + steps: steps.clone(), + is_running: is_running, + session_findings: findings_count as usize, + session_tool_invocations: tool_invocations as usize, + session_success_rate: success_rate, + } } + } + }, + Some(None) => rsx! { p { style: "color: var(--text-secondary);", "Failed to load attack chain." } }, + None => rsx! { p { style: "color: var(--text-secondary);", "Loading..." } }, + } + } + } + } + + // Export modal + if *show_export_modal.read() { + div { + style: "position: fixed; inset: 0; background: rgba(0,0,0,0.6); display: flex; align-items: center; justify-content: center; z-index: 1000;", + onclick: move |_| show_export_modal.set(false), + div { + style: "background: var(--bg-secondary); border: 1px solid var(--border-color); border-radius: 12px; padding: 24px; width: 480px; max-width: 90vw;", + onclick: move |e| e.stop_propagation(), + h3 { style: "margin: 0 0 4px 0;", "Export Pentest Report" } + p { style: "font-size: 0.85rem; color: var(--text-secondary); margin-bottom: 16px;", + "The report will be exported as a password-protected ZIP archive (AES-256) containing a professional HTML report and raw findings data. Open with any standard archive tool." + } + + div { style: "margin-bottom: 14px;", + label { style: "display: block; font-size: 0.85rem; color: var(--text-secondary); margin-bottom: 4px;", + "Encryption Password" + } + input { + class: "chat-input", + style: "width: 100%; padding: 8px;", + r#type: "password", + placeholder: "Minimum 8 characters", + value: "{export_password}", + oninput: move |e| { + export_password.set(e.value()); + export_error.set(None); + }, + } + } + + if let Some(err) = &*export_error.read() { + div { style: "padding: 8px 12px; background: rgba(220, 38, 38, 0.1); border: 1px solid #dc2626; border-radius: 6px; color: #dc2626; font-size: 0.85rem; margin-bottom: 14px;", + "{err}" + } + } + + if let Some(sha) = &*export_sha256.read() { + { + let sha_copy = sha.clone(); + rsx! { + div { style: "padding: 10px 12px; background: rgba(22, 163, 74, 0.08); border: 1px solid #16a34a; border-radius: 6px; margin-bottom: 14px;", + div { style: "font-size: 0.8rem; font-weight: 600; color: #16a34a; margin-bottom: 4px;", + Icon { icon: BsCheckCircle, width: 12, height: 12 } + " Archive downloaded successfully" + } + div { style: "font-size: 0.75rem; color: var(--text-secondary); margin-bottom: 2px;", + "SHA-256 Checksum:" + } + div { style: "display: flex; align-items: center; gap: 6px;", + div { style: "flex: 1; font-family: monospace; font-size: 0.7rem; word-break: break-all; color: var(--text-primary); background: var(--bg-primary); padding: 6px 8px; border-radius: 4px;", + "{sha_copy}" + } + button { + class: "btn btn-ghost", + style: "padding: 4px 8px; font-size: 0.75rem; flex-shrink: 0;", + onclick: move |_| { + let js = format!( + "navigator.clipboard.writeText('{}');", + sha_copy + ); + document::eval(&js); + }, + Icon { icon: BsClipboard, width: 12, height: 12 } + } + } + } + } + } + } + + div { style: "display: flex; justify-content: flex-end; gap: 8px;", + button { + class: "btn btn-ghost", + onclick: move |_| show_export_modal.set(false), + "Close" + } + button { + class: "btn btn-primary", + disabled: *exporting.read() || export_password.read().len() < 8, + onclick: do_export, + if *exporting.read() { "Encrypting..." } else { "Export" } + } + } + } + } + } + } +} + +// ═══════════════════════════════════════ +// Attack Chain Visualization Component +// ═══════════════════════════════════════ + +/// Get category CSS class from tool name +fn tool_category(name: &str) -> &'static str { + let lower = name.to_lowercase(); + if lower.contains("recon") { return "recon"; } + if lower.contains("openapi") || lower.contains("api") || lower.contains("swagger") { return "api"; } + if lower.contains("header") { return "headers"; } + if lower.contains("csp") { return "csp"; } + if lower.contains("cookie") { return "cookies"; } + if lower.contains("log") || lower.contains("console") { return "logs"; } + if lower.contains("rate") || lower.contains("limit") { return "ratelimit"; } + if lower.contains("cors") { return "cors"; } + if lower.contains("tls") || lower.contains("ssl") { return "tls"; } + if lower.contains("redirect") { return "redirect"; } + if lower.contains("dns") || lower.contains("dmarc") || lower.contains("email") || lower.contains("spf") { return "email"; } + if lower.contains("auth") || lower.contains("jwt") || lower.contains("token") || lower.contains("session") { return "auth"; } + if lower.contains("xss") { return "xss"; } + if lower.contains("sql") || lower.contains("sqli") { return "sqli"; } + if lower.contains("ssrf") { return "ssrf"; } + if lower.contains("idor") { return "idor"; } + if lower.contains("fuzz") { return "fuzzer"; } + if lower.contains("cve") || lower.contains("exploit") { return "cve"; } + "default" +} + +/// Get emoji icon from tool category +fn tool_emoji(cat: &str) -> &'static str { + match cat { + "recon" => "\u{1F50D}", + "api" => "\u{1F517}", + "headers" => "\u{1F6E1}", + "csp" => "\u{1F6A7}", + "cookies" => "\u{1F36A}", + "logs" => "\u{1F4DD}", + "ratelimit" => "\u{23F1}", + "cors" => "\u{1F30D}", + "tls" => "\u{1F510}", + "redirect" => "\u{21AA}", + "email" => "\u{1F4E7}", + "auth" => "\u{1F512}", + "xss" => "\u{26A1}", + "sqli" => "\u{1F489}", + "ssrf" => "\u{1F310}", + "idor" => "\u{1F511}", + "fuzzer" => "\u{1F9EA}", + "cve" => "\u{1F4A3}", + _ => "\u{1F527}", + } +} + +/// Compute display label for category +fn cat_label(cat: &str) -> &'static str { + match cat { + "recon" => "Recon", + "api" => "API", + "headers" => "Headers", + "csp" => "CSP", + "cookies" => "Cookies", + "logs" => "Logs", + "ratelimit" => "Rate Limit", + "cors" => "CORS", + "tls" => "TLS", + "redirect" => "Redirect", + "email" => "Email/DNS", + "auth" => "Auth", + "xss" => "XSS", + "sqli" => "SQLi", + "ssrf" => "SSRF", + "idor" => "IDOR", + "fuzzer" => "Fuzzer", + "cve" => "CVE", + _ => "Other", + } +} + +/// Phase name heuristic based on depth +fn phase_name(depth: usize) -> &'static str { + match depth { + 0 => "Reconnaissance", + 1 => "Analysis", + 2 => "Boundary Testing", + 3 => "Injection & Exploitation", + 4 => "Authentication Testing", + 5 => "Validation", + 6 => "Deep Scan", + _ => "Final", + } +} + +/// Short label for phase rail +fn phase_short_name(depth: usize) -> &'static str { + match depth { + 0 => "Recon", + 1 => "Analysis", + 2 => "Boundary", + 3 => "Exploit", + 4 => "Auth", + 5 => "Validate", + 6 => "Deep", + _ => "Final", + } +} + +/// Compute BFS phases from attack chain nodes +fn compute_phases(steps: &[serde_json::Value]) -> Vec> { + let node_ids: Vec = steps + .iter() + .map(|s| s.get("node_id").and_then(|v| v.as_str()).unwrap_or("").to_string()) + .collect(); + + let id_to_idx: HashMap = node_ids + .iter() + .enumerate() + .map(|(i, id)| (id.clone(), i)) + .collect(); + + // Compute depth via BFS + let mut depths = vec![usize::MAX; steps.len()]; + let mut queue = VecDeque::new(); + + // Root nodes: those with no parents or parents not in the set + for (i, step) in steps.iter().enumerate() { + let parents = step + .get("parent_node_ids") + .and_then(|v| v.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|p| p.as_str()) + .filter(|p| id_to_idx.contains_key(*p)) + .count() + }) + .unwrap_or(0); + if parents == 0 { + depths[i] = 0; + queue.push_back(i); + } + } + + // BFS to compute min depth + while let Some(idx) = queue.pop_front() { + let current_depth = depths[idx]; + let node_id = &node_ids[idx]; + // Find children: nodes that list this node as a parent + for (j, step) in steps.iter().enumerate() { + if depths[j] <= current_depth + 1 { + continue; + } + let is_child = step + .get("parent_node_ids") + .and_then(|v| v.as_array()) + .map(|arr| arr.iter().any(|p| p.as_str() == Some(node_id.as_str()))) + .unwrap_or(false); + if is_child { + depths[j] = current_depth + 1; + queue.push_back(j); + } + } + } + + // Handle unreachable nodes + for d in depths.iter_mut() { + if *d == usize::MAX { + *d = 0; + } + } + + // Group by depth + let max_depth = depths.iter().copied().max().unwrap_or(0); + let mut phases: Vec> = Vec::new(); + for d in 0..=max_depth { + let indices: Vec = depths + .iter() + .enumerate() + .filter(|(_, &dep)| dep == d) + .map(|(i, _)| i) + .collect(); + if !indices.is_empty() { + phases.push(indices); + } + } + phases +} + +/// Format BSON datetime to readable string +fn format_bson_time(val: &serde_json::Value) -> String { + // Handle BSON {"$date":{"$numberLong":"..."}} + if let Some(date_obj) = val.get("$date") { + if let Some(ms_str) = date_obj.get("$numberLong").and_then(|v| v.as_str()) { + if let Ok(ms) = ms_str.parse::() { + let secs = ms / 1000; + let h = (secs / 3600) % 24; + let m = (secs / 60) % 60; + let s = secs % 60; + return format!("{h:02}:{m:02}:{s:02}"); + } + } + // Handle {"$date": "2025-..."} + if let Some(s) = date_obj.as_str() { + return s.to_string(); + } + } + // Handle plain string + if let Some(s) = val.as_str() { + return s.to_string(); + } + String::new() +} + +/// Compute duration string from started_at and completed_at +fn compute_duration(step: &serde_json::Value) -> String { + let extract_ms = |val: &serde_json::Value| -> Option { + val.get("$date")? + .get("$numberLong")? + .as_str()? + .parse::() + .ok() + }; + + let started = step.get("started_at").and_then(extract_ms); + let completed = step.get("completed_at").and_then(extract_ms); + + match (started, completed) { + (Some(s), Some(c)) => { + let diff_ms = c - s; + if diff_ms < 1000 { + format!("{}ms", diff_ms) + } else { + format!("{:.1}s", diff_ms as f64 / 1000.0) + } + } + _ => String::new(), + } +} + +#[component] +fn AttackChainView( + steps: Vec, + is_running: bool, + session_findings: usize, + session_tool_invocations: usize, + session_success_rate: f64, +) -> Element { + let phases = compute_phases(&steps); + + // Compute KPIs — prefer session-level stats, fall back to node-level + let total_tools = steps.len(); + let node_findings: usize = steps + .iter() + .map(|s| { + s.get("findings_produced") + .and_then(|v| v.as_array()) + .map(|a| a.len()) + .unwrap_or(0) + }) + .sum(); + // Use session-level findings count if nodes don't have findings linked + let total_findings = if node_findings > 0 { node_findings } else { session_findings }; + + let completed_count = steps + .iter() + .filter(|s| s.get("status").and_then(|v| v.as_str()) == Some("completed")) + .count(); + let failed_count = steps + .iter() + .filter(|s| s.get("status").and_then(|v| v.as_str()) == Some("failed")) + .count(); + let finished = completed_count + failed_count; + let success_pct = if finished == 0 { + 100 + } else { + (completed_count * 100) / finished + }; + let max_risk: u8 = steps + .iter() + .filter_map(|s| s.get("risk_score").and_then(|v| v.as_u64())) + .map(|v| v as u8) + .max() + .unwrap_or(0); + + let progress_pct = if total_tools == 0 { + 0 + } else { + ((completed_count + failed_count) * 100) / total_tools + }; + + // Build phase data for rail and accordion + let phase_data: Vec<(usize, Vec<&serde_json::Value>, usize, bool, bool, bool)> = phases + .iter() + .enumerate() + .map(|(pi, indices)| { + let phase_steps: Vec<&serde_json::Value> = indices.iter().map(|&i| &steps[i]).collect(); + let phase_findings: usize = phase_steps + .iter() + .map(|s| { + s.get("findings_produced") + .and_then(|v| v.as_array()) + .map(|a| a.len()) + .unwrap_or(0) + }) + .sum(); + let has_failed = phase_steps + .iter() + .any(|s| s.get("status").and_then(|v| v.as_str()) == Some("failed")); + let has_running = phase_steps + .iter() + .any(|s| s.get("status").and_then(|v| v.as_str()) == Some("running")); + let all_done = phase_steps.iter().all(|s| { + let st = s.get("status").and_then(|v| v.as_str()).unwrap_or(""); + st == "completed" || st == "failed" || st == "skipped" + }); + (pi, phase_steps, phase_findings, has_failed, has_running, all_done) + }) + .collect(); + + let mut active_rail = use_signal(|| 0usize); + + rsx! { + // KPI bar + div { class: "ac-kpi-bar", + div { class: "ac-kpi-card", + div { class: "ac-kpi-value", style: "color: var(--text-primary);", "{total_tools}" } + div { class: "ac-kpi-label", "Tools Run" } + } + div { class: "ac-kpi-card", + div { class: "ac-kpi-value", style: "color: var(--danger, #dc2626);", "{total_findings}" } + div { class: "ac-kpi-label", "Findings" } + } + div { class: "ac-kpi-card", + div { class: "ac-kpi-value", style: "color: var(--success, #16a34a);", "{success_pct}%" } + div { class: "ac-kpi-label", "Success Rate" } + } + div { class: "ac-kpi-card", + div { class: "ac-kpi-value", style: "color: var(--warning, #d97706);", "{max_risk}" } + div { class: "ac-kpi-label", "Max Risk" } + } + } + + // Phase rail + div { class: "ac-phase-rail", + for (pi, (_phase_idx, phase_steps, phase_findings, has_failed, has_running, all_done)) in phase_data.iter().enumerate() { + { + if pi > 0 { + let prev_done = phase_data.get(pi - 1).map(|p| p.5).unwrap_or(false); + let bar_class = if prev_done && *all_done { + "done" + } else if prev_done { + "running" + } else { + "" + }; + rsx! { + div { class: "ac-rail-bar", + div { class: "ac-rail-bar-inner {bar_class}" } + } + } + } else { + rsx! {} + } + } + { + let dot_class = if *has_running { + "running" + } else if *has_failed && *all_done { + "mixed" + } else if *all_done { + "done" + } else { + "pending" + }; + let is_active = *active_rail.read() == pi; + let active_cls = if is_active { " active" } else { "" }; + let findings_cls = if *phase_findings > 0 { "has" } else { "none" }; + let findings_text = if *phase_findings > 0 { + format!("{phase_findings}") + } else { + "\u{2014}".to_string() + }; + let short = phase_short_name(pi); + + rsx! { + div { + class: "ac-rail-node{active_cls}", + onclick: move |_| { + active_rail.set(pi); + let js = format!( + "document.getElementById('ac-phase-{pi}')?.scrollIntoView({{behavior:'smooth',block:'nearest'}});document.getElementById('ac-phase-{pi}')?.classList.add('open');" + ); + document::eval(&js); + }, + div { class: "ac-rail-dot {dot_class}" } + div { class: "ac-rail-label", "{short}" } + div { class: "ac-rail-findings {findings_cls}", "{findings_text}" } + div { class: "ac-rail-heatmap", + for step in phase_steps.iter() { + { + let st = step.get("status").and_then(|v| v.as_str()).unwrap_or("pending"); + let hm_cls = match st { + "completed" => "ok", + "failed" => "fail", + "running" => "run", + _ => "wait", + }; + rsx! { div { class: "ac-hm-cell {hm_cls}" } } + } + } + } + } + } + } + } + } + + // Progress bar + div { class: "ac-progress-track", + div { class: "ac-progress-fill", style: "width: {progress_pct}%;" } + } + + // Expand all + div { class: "ac-controls", + button { + class: "ac-btn-toggle", + onclick: move |_| { + document::eval( + "document.querySelectorAll('.ac-phase').forEach(p => p.classList.toggle('open', !document.querySelector('.ac-phase.open') || !document.querySelectorAll('.ac-phase:not(.open)').length === 0));(function(){var ps=document.querySelectorAll('.ac-phase');var allOpen=Array.from(ps).every(p=>p.classList.contains('open'));ps.forEach(p=>{if(allOpen)p.classList.remove('open');else p.classList.add('open');});})();" + ); + }, + "Expand all" + } + } + + // Phase accordion + div { class: "ac-phases", + for (pi, (_, phase_steps, phase_findings, has_failed, has_running, all_done)) in phase_data.iter().enumerate() { + { + let open_cls = if pi == 0 { " open" } else { "" }; + let phase_label = phase_name(pi); + let tool_count = phase_steps.len(); + let meta_text = if *has_running { + "in progress".to_string() + } else { + format!("{phase_findings} findings") + }; + let meta_cls = if *has_running { "running-ct" } else { "findings-ct" }; + let phase_num_label = format!("PHASE {}", pi + 1); + let phase_el_id = format!("ac-phase-{pi}"); + let phase_el_id2 = phase_el_id.clone(); + + rsx! { + div { + class: "ac-phase{open_cls}", + id: "{phase_el_id}", + div { + class: "ac-phase-header", + onclick: move |_| { + let js = format!("document.getElementById('{phase_el_id2}').classList.toggle('open');"); + document::eval(&js); + }, + span { class: "ac-phase-num", "{phase_num_label}" } + span { class: "ac-phase-title", "{phase_label}" } + div { class: "ac-phase-dots", + for step in phase_steps.iter() { + { + let st = step.get("status").and_then(|v| v.as_str()).unwrap_or("pending"); + rsx! { div { class: "ac-phase-dot {st}" } } + } + } + } + div { class: "ac-phase-meta", + span { "{tool_count} tools" } + span { class: "{meta_cls}", "{meta_text}" } + } + span { class: "ac-phase-chevron", "\u{25B8}" } + } + div { class: "ac-phase-body", + div { class: "ac-phase-body-inner", + for step in phase_steps.iter() { + { + let tool_name_val = step.get("tool_name").and_then(|v| v.as_str()).unwrap_or("Unknown").to_string(); + let status = step.get("status").and_then(|v| v.as_str()).unwrap_or("pending").to_string(); + let cat = tool_category(&tool_name_val); + let emoji = tool_emoji(cat); + let label = cat_label(cat); + let findings_n = step.get("findings_produced").and_then(|v| v.as_array()).map(|a| a.len()).unwrap_or(0); + let risk = step.get("risk_score").and_then(|v| v.as_u64()).map(|v| v as u8); + let reasoning = step.get("llm_reasoning").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let duration = compute_duration(step); + let started = step.get("started_at").map(format_bson_time).unwrap_or_default(); + + let is_pending = status == "pending"; + let pending_cls = if is_pending { " is-pending" } else { "" }; + + let duration_cls = if status == "running" { "ac-tool-duration running-text" } else { "ac-tool-duration" }; + let duration_text = if status == "running" { + "running\u{2026}".to_string() + } else if duration.is_empty() { + "\u{2014}".to_string() + } else { + duration + }; + + let pill_cls = if findings_n > 0 { "ac-findings-pill has" } else { "ac-findings-pill zero" }; + let pill_text = if findings_n > 0 { format!("{findings_n}") } else { "\u{2014}".to_string() }; + + let (risk_cls, risk_text) = match risk { + Some(r) if r >= 75 => ("ac-risk-val high", format!("{r}")), + Some(r) if r >= 40 => ("ac-risk-val medium", format!("{r}")), + Some(r) => ("ac-risk-val low", format!("{r}")), + None => ("ac-risk-val none", "\u{2014}".to_string()), + }; + + let node_id = step.get("node_id").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let detail_id = format!("ac-detail-{node_id}"); + let row_id = format!("ac-row-{node_id}"); + let detail_id_clone = detail_id.clone(); + + rsx! { + div { + class: "ac-tool-row{pending_cls}", + id: "{row_id}", + onclick: move |_| { + if is_pending { return; } + let js = format!( + "(function(){{var r=document.getElementById('{row_id}');var d=document.getElementById('{detail_id}');if(r.classList.contains('expanded')){{r.classList.remove('expanded');d.classList.remove('open');}}else{{r.classList.add('expanded');d.classList.add('open');}}}})()" + ); + document::eval(&js); + }, + div { class: "ac-status-bar {status}" } + div { class: "ac-tool-icon", "{emoji}" } + div { class: "ac-tool-info", + div { class: "ac-tool-name", "{tool_name_val}" } + span { class: "ac-cat-chip {cat}", "{label}" } + } + div { class: "{duration_cls}", "{duration_text}" } + div { span { class: "{pill_cls}", "{pill_text}" } } + div { class: "{risk_cls}", "{risk_text}" } + } + div { + class: "ac-tool-detail", + id: "{detail_id_clone}", + if !reasoning.is_empty() || !started.is_empty() { + div { class: "ac-tool-detail-inner", + if !reasoning.is_empty() { + div { class: "ac-reasoning-block", "{reasoning}" } + } + if !started.is_empty() { + div { class: "ac-detail-grid", + span { class: "ac-detail-label", "Started" } + span { class: "ac-detail-value", "{started}" } + if !duration_text.is_empty() && status != "running" && duration_text != "\u{2014}" { + span { class: "ac-detail-label", "Duration" } + span { class: "ac-detail-value", "{duration_text}" } + } + span { class: "ac-detail-label", "Status" } + if status == "completed" { + span { class: "ac-detail-value", style: "color: var(--success, #16a34a);", "Completed" } + } else if status == "failed" { + span { class: "ac-detail-value", style: "color: var(--danger, #dc2626);", "Failed" } + } else if status == "running" { + span { class: "ac-detail-value", style: "color: var(--warning, #d97706);", "Running" } + } else { + span { class: "ac-detail-value", "{status}" } + } + } + } + } + } + } + } + } + } + } + } + } + } + } + } + } + } +} diff --git a/compliance-dast/Cargo.toml b/compliance-dast/Cargo.toml index 4613e24..e91ef99 100644 --- a/compliance-dast/Cargo.toml +++ b/compliance-dast/Cargo.toml @@ -27,6 +27,10 @@ chromiumoxide = { version = "0.7", features = ["tokio-runtime"], default-feature # Docker sandboxing bollard = "0.18" +# TLS analysis +native-tls = "0.2" +tokio-native-tls = "0.3" + # Serialization bson = { version = "2", features = ["chrono-0_4"] } url = "2" diff --git a/compliance-dast/src/lib.rs b/compliance-dast/src/lib.rs index 38c2445..aefc932 100644 --- a/compliance-dast/src/lib.rs +++ b/compliance-dast/src/lib.rs @@ -2,5 +2,7 @@ pub mod agents; pub mod crawler; pub mod orchestrator; pub mod recon; +pub mod tools; pub use orchestrator::DastOrchestrator; +pub use tools::ToolRegistry; diff --git a/compliance-dast/src/tools/api_fuzzer.rs b/compliance-dast/src/tools/api_fuzzer.rs new file mode 100644 index 0000000..39d3f7e --- /dev/null +++ b/compliance-dast/src/tools/api_fuzzer.rs @@ -0,0 +1,146 @@ +use compliance_core::error::CoreError; +use compliance_core::traits::dast_agent::{DastAgent, DastContext, DiscoveredEndpoint, EndpointParameter}; +use compliance_core::traits::pentest_tool::{PentestTool, PentestToolContext, PentestToolResult}; +use serde_json::json; + +use crate::agents::api_fuzzer::ApiFuzzerAgent; + +/// PentestTool wrapper around the existing ApiFuzzerAgent. +pub struct ApiFuzzerTool { + http: reqwest::Client, + agent: ApiFuzzerAgent, +} + +impl ApiFuzzerTool { + pub fn new(http: reqwest::Client) -> Self { + let agent = ApiFuzzerAgent::new(http.clone()); + Self { http, agent } + } + + fn parse_endpoints(input: &serde_json::Value) -> Vec { + let mut endpoints = Vec::new(); + if let Some(arr) = input.get("endpoints").and_then(|v| v.as_array()) { + for ep in arr { + let url = ep.get("url").and_then(|v| v.as_str()).unwrap_or_default().to_string(); + let method = ep.get("method").and_then(|v| v.as_str()).unwrap_or("GET").to_string(); + let mut parameters = Vec::new(); + if let Some(params) = ep.get("parameters").and_then(|v| v.as_array()) { + for p in params { + parameters.push(EndpointParameter { + name: p.get("name").and_then(|v| v.as_str()).unwrap_or_default().to_string(), + location: p.get("location").and_then(|v| v.as_str()).unwrap_or("query").to_string(), + param_type: p.get("param_type").and_then(|v| v.as_str()).map(String::from), + example_value: p.get("example_value").and_then(|v| v.as_str()).map(String::from), + }); + } + } + endpoints.push(DiscoveredEndpoint { + url, + method, + parameters, + content_type: ep.get("content_type").and_then(|v| v.as_str()).map(String::from), + requires_auth: ep.get("requires_auth").and_then(|v| v.as_bool()).unwrap_or(false), + }); + } + } + endpoints + } +} + +impl PentestTool for ApiFuzzerTool { + fn name(&self) -> &str { + "api_fuzzer" + } + + fn description(&self) -> &str { + "Fuzzes API endpoints to discover misconfigurations, information disclosure, and hidden \ + endpoints. Probes common sensitive paths and tests for verbose error messages." + } + + fn input_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "endpoints": { + "type": "array", + "description": "Known endpoints to fuzz", + "items": { + "type": "object", + "properties": { + "url": { "type": "string" }, + "method": { "type": "string", "enum": ["GET", "POST", "PUT", "PATCH", "DELETE"] }, + "parameters": { + "type": "array", + "items": { + "type": "object", + "properties": { + "name": { "type": "string" }, + "location": { "type": "string" }, + "param_type": { "type": "string" }, + "example_value": { "type": "string" } + }, + "required": ["name"] + } + } + }, + "required": ["url"] + } + }, + "base_url": { + "type": "string", + "description": "Base URL to probe for common sensitive paths (used if no endpoints provided)" + } + } + }) + } + + fn execute<'a>( + &'a self, + input: serde_json::Value, + context: &'a PentestToolContext, + ) -> std::pin::Pin> + Send + 'a>> { + Box::pin(async move { + let mut endpoints = Self::parse_endpoints(&input); + + // If a base_url is provided but no endpoints, create a default endpoint + if endpoints.is_empty() { + if let Some(base) = input.get("base_url").and_then(|v| v.as_str()) { + endpoints.push(DiscoveredEndpoint { + url: base.to_string(), + method: "GET".to_string(), + parameters: Vec::new(), + content_type: None, + requires_auth: false, + }); + } + } + + if endpoints.is_empty() { + return Ok(PentestToolResult { + summary: "No endpoints or base_url provided to fuzz.".to_string(), + findings: Vec::new(), + data: json!({}), + }); + } + + let dast_context = DastContext { + endpoints, + technologies: Vec::new(), + sast_hints: Vec::new(), + }; + + let findings = self.agent.run(&context.target, &dast_context).await?; + let count = findings.len(); + + Ok(PentestToolResult { + summary: if count > 0 { + format!("Found {count} API misconfigurations or information disclosures.") + } else { + "No API misconfigurations detected.".to_string() + }, + findings, + data: json!({ "endpoints_tested": dast_context.endpoints.len() }), + }) + }) + } +} diff --git a/compliance-dast/src/tools/auth_bypass.rs b/compliance-dast/src/tools/auth_bypass.rs new file mode 100644 index 0000000..6f0bf87 --- /dev/null +++ b/compliance-dast/src/tools/auth_bypass.rs @@ -0,0 +1,130 @@ +use compliance_core::error::CoreError; +use compliance_core::traits::dast_agent::{DastAgent, DastContext, DiscoveredEndpoint, EndpointParameter}; +use compliance_core::traits::pentest_tool::{PentestTool, PentestToolContext, PentestToolResult}; +use serde_json::json; + +use crate::agents::auth_bypass::AuthBypassAgent; + +/// PentestTool wrapper around the existing AuthBypassAgent. +pub struct AuthBypassTool { + http: reqwest::Client, + agent: AuthBypassAgent, +} + +impl AuthBypassTool { + pub fn new(http: reqwest::Client) -> Self { + let agent = AuthBypassAgent::new(http.clone()); + Self { http, agent } + } + + fn parse_endpoints(input: &serde_json::Value) -> Vec { + let mut endpoints = Vec::new(); + if let Some(arr) = input.get("endpoints").and_then(|v| v.as_array()) { + for ep in arr { + let url = ep.get("url").and_then(|v| v.as_str()).unwrap_or_default().to_string(); + let method = ep.get("method").and_then(|v| v.as_str()).unwrap_or("GET").to_string(); + let mut parameters = Vec::new(); + if let Some(params) = ep.get("parameters").and_then(|v| v.as_array()) { + for p in params { + parameters.push(EndpointParameter { + name: p.get("name").and_then(|v| v.as_str()).unwrap_or_default().to_string(), + location: p.get("location").and_then(|v| v.as_str()).unwrap_or("query").to_string(), + param_type: p.get("param_type").and_then(|v| v.as_str()).map(String::from), + example_value: p.get("example_value").and_then(|v| v.as_str()).map(String::from), + }); + } + } + endpoints.push(DiscoveredEndpoint { + url, + method, + parameters, + content_type: ep.get("content_type").and_then(|v| v.as_str()).map(String::from), + requires_auth: ep.get("requires_auth").and_then(|v| v.as_bool()).unwrap_or(false), + }); + } + } + endpoints + } +} + +impl PentestTool for AuthBypassTool { + fn name(&self) -> &str { + "auth_bypass_scanner" + } + + fn description(&self) -> &str { + "Tests endpoints for authentication bypass vulnerabilities. Tries accessing protected \ + endpoints without credentials, with manipulated tokens, and with common default credentials." + } + + fn input_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "endpoints": { + "type": "array", + "description": "Endpoints to test for authentication bypass", + "items": { + "type": "object", + "properties": { + "url": { "type": "string" }, + "method": { "type": "string", "enum": ["GET", "POST", "PUT", "PATCH", "DELETE"] }, + "parameters": { + "type": "array", + "items": { + "type": "object", + "properties": { + "name": { "type": "string" }, + "location": { "type": "string" }, + "param_type": { "type": "string" }, + "example_value": { "type": "string" } + }, + "required": ["name"] + } + }, + "requires_auth": { "type": "boolean", "description": "Whether this endpoint requires authentication" } + }, + "required": ["url", "method"] + } + } + }, + "required": ["endpoints"] + }) + } + + fn execute<'a>( + &'a self, + input: serde_json::Value, + context: &'a PentestToolContext, + ) -> std::pin::Pin> + Send + 'a>> { + Box::pin(async move { + let endpoints = Self::parse_endpoints(&input); + if endpoints.is_empty() { + return Ok(PentestToolResult { + summary: "No endpoints provided to test.".to_string(), + findings: Vec::new(), + data: json!({}), + }); + } + + let dast_context = DastContext { + endpoints, + technologies: Vec::new(), + sast_hints: Vec::new(), + }; + + let findings = self.agent.run(&context.target, &dast_context).await?; + let count = findings.len(); + + Ok(PentestToolResult { + summary: if count > 0 { + format!("Found {count} authentication bypass vulnerabilities.") + } else { + "No authentication bypass vulnerabilities detected.".to_string() + }, + findings, + data: json!({ "endpoints_tested": dast_context.endpoints.len() }), + }) + }) + } +} diff --git a/compliance-dast/src/tools/console_log_detector.rs b/compliance-dast/src/tools/console_log_detector.rs new file mode 100644 index 0000000..c1bc9cd --- /dev/null +++ b/compliance-dast/src/tools/console_log_detector.rs @@ -0,0 +1,326 @@ +use compliance_core::error::CoreError; +use compliance_core::models::dast::{DastEvidence, DastFinding, DastVulnType}; +use compliance_core::models::Severity; +use compliance_core::traits::pentest_tool::{PentestTool, PentestToolContext, PentestToolResult}; +use serde_json::json; +use tracing::info; + +/// Tool that detects console.log and similar debug statements in frontend JavaScript. +pub struct ConsoleLogDetectorTool { + http: reqwest::Client, +} + +/// A detected console statement with its context. +#[derive(Debug)] +struct ConsoleMatch { + pattern: String, + file_url: String, + line_snippet: String, + line_number: Option, +} + +impl ConsoleLogDetectorTool { + pub fn new(http: reqwest::Client) -> Self { + Self { http } + } + + /// Patterns that indicate debug/logging statements left in production code. + fn patterns() -> Vec<&'static str> { + vec![ + "console.log(", + "console.debug(", + "console.error(", + "console.warn(", + "console.info(", + "console.trace(", + "console.dir(", + "console.table(", + "debugger;", + "alert(", + ] + } + + /// Extract JavaScript file URLs from an HTML page body. + fn extract_js_urls(html: &str, base_url: &str) -> Vec { + let mut urls = Vec::new(); + let base = url::Url::parse(base_url).ok(); + + // Simple regex-free extraction of