Some checks failed
CI / Check (push) Has been skipped
CI / Deploy Agent (push) Has been cancelled
CI / Deploy Dashboard (push) Has been cancelled
CI / Deploy Docs (push) Has been cancelled
CI / Deploy MCP (push) Has been cancelled
CI / Detect Changes (push) Has been cancelled
42 lines
2.5 KiB
Rust
42 lines
2.5 KiB
Rust
use std::sync::Arc;
|
|
|
|
use compliance_core::models::Finding;
|
|
|
|
use crate::error::AgentError;
|
|
use crate::llm::LlmClient;
|
|
|
|
const FIX_SYSTEM_PROMPT: &str = r#"You are a security engineer suggesting a code fix. Return ONLY the corrected code that replaces the vulnerable snippet — no explanations, no markdown fences, no before/after comparison.
|
|
|
|
Rules:
|
|
- The fix must be a drop-in replacement for the vulnerable code
|
|
- Preserve the original code's style, indentation, and naming conventions
|
|
- Add at most one brief inline comment on the changed line explaining the security fix
|
|
- If the fix requires importing a new module, include the import on a separate line prefixed with the language's comment syntax + "Add import: "
|
|
- Do not refactor, rename variables, or "improve" unrelated code
|
|
- If the vulnerability is a false positive and the code is actually safe, return the original code unchanged with a comment explaining why no fix is needed
|
|
|
|
Language-specific fix guidance:
|
|
- Rust: use `?` for error propagation, prefer `SecretString` for secrets, use parameterized queries with `sqlx`/`diesel`
|
|
- Python: use parameterized queries (never f-strings in SQL), use `secrets` module not `random`, use `subprocess.run([...])` list form, use `markupsafe.escape()` for HTML
|
|
- Go: use `sql.Query` with `$1`/`?` placeholders, use `crypto/rand` not `math/rand`, use `html/template` not `text/template`, return errors don't panic
|
|
- Java/Kotlin: use `PreparedStatement` with `?` params, use `SecureRandom`, use `Jsoup.clean()` for HTML sanitization, use `@Valid` for input validation
|
|
- Ruby: use ActiveRecord parameterized finders, use `SecureRandom`, use `ERB::Util.html_escape`, use `strong_parameters`
|
|
- PHP: use PDO prepared statements with `:param` or `?`, use `random_bytes()`/`random_int()`, use `htmlspecialchars()` with `ENT_QUOTES`, use `password_hash(PASSWORD_BCRYPT)`
|
|
- C/C++: use `snprintf` not `sprintf`, use bounds-checked APIs, free resources in reverse allocation order, use `memset_s` for secret cleanup"#;
|
|
|
|
pub async fn suggest_fix(llm: &Arc<LlmClient>, finding: &Finding) -> Result<String, AgentError> {
|
|
let user_prompt = format!(
|
|
"Suggest a fix for this vulnerability:\n\
|
|
Language context from file: {}\n\
|
|
Rule: {}\n\
|
|
Description: {}\n\
|
|
Vulnerable code:\n```\n{}\n```",
|
|
finding.file_path.as_deref().unwrap_or("unknown"),
|
|
finding.rule_id.as_deref().unwrap_or("N/A"),
|
|
finding.description,
|
|
finding.code_snippet.as_deref().unwrap_or("N/A"),
|
|
);
|
|
|
|
llm.chat(FIX_SYSTEM_PROMPT, &user_prompt, Some(0.2)).await
|
|
}
|