repl scan feature

This commit is contained in:
rarebuffalo
2026-05-22 21:47:45 +05:30
parent 5826119812
commit d94757eee2

View File

@@ -5,35 +5,36 @@ Post-scan Q&A loop — the "Gemini CLI feel".
After a scan completes, the user drops into this loop where they can: After a scan completes, the user drops into this loop where they can:
- Ask natural-language questions about the scan results - Ask natural-language questions about the scan results
- Use slash commands (/export, /files, /model, /clear, /help) - Use slash commands (/export, /files, /score, /model, /clear, /help, /exit)
- Ctrl+C to exit - Ctrl+C to exit
The AI is given full scan context at the start of the conversation The AI is given full scan context at the start of the conversation
and remembers the entire chat history during the session. and remembers the entire chat history during the session.
""" """
import asyncio
import json import json
import sys
from dataclasses import dataclass, field from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional from typing import Optional
from rich.console import Console from rich.console import Console
from rich.markdown import Markdown
from rich.panel import Panel
from rich.prompt import Prompt
from securelens.ai import call_ai from securelens.ai import call_ai
from securelens.ai.prompts import chat_prompt from securelens.ai.prompts import chat_prompt
from securelens.output import console, print_ai_response, print_info, print_error, print_success
from securelens.output.exporters import save_json, save_markdown from securelens.output.exporters import save_json, save_markdown
console_out = Console() console = Console()
HELP_TEXT = """ HELP_TEXT = """
[bold cyan]Available commands:[/bold cyan] [bold cyan]Available commands:[/bold cyan]
[bold]/help[/bold] Show this help message [bold]/help[/bold] Show this help message
[bold]/files[/bold] List files that were scanned [bold]/files[/bold] List files that were analyzed
[bold]/score[/bold] Show the security score [bold]/score[/bold] Show the current security score
[bold]/issues[/bold] Show all found issues (summary)
[bold]/issues critical[/bold] Filter issues by severity (critical/high/medium/low)
[bold]/export markdown[/bold] Save the report as a Markdown file [bold]/export markdown[/bold] Save the report as a Markdown file
[bold]/export json[/bold] Save the report as a JSON file [bold]/export json[/bold] Save the report as a JSON file
[bold]/model <name>[/bold] Switch AI model (e.g. /model gpt-4o-mini) [bold]/model <name>[/bold] Switch AI model (e.g. /model gpt-4o-mini)
@@ -44,6 +45,7 @@ Or just type a question in plain English, e.g.:
[dim]> How do I fix the SQL injection?[/dim] [dim]> How do I fix the SQL injection?[/dim]
[dim]> What's the most critical issue?[/dim] [dim]> What's the most critical issue?[/dim]
[dim]> Show me all issues in auth.py[/dim] [dim]> Show me all issues in auth.py[/dim]
[dim]> Give me a step-by-step remediation plan[/dim]
""" """
@@ -51,7 +53,7 @@ Or just type a question in plain English, e.g.:
class ReplContext: class ReplContext:
target: str target: str
scan_result: object # LocalScanResult or WebScanResult scan_result: object # LocalScanResult or WebScanResult
target_type: str # "code" | "web" target_type: str # "code" | "web" | "github"
api_key: str api_key: str
model: str model: str
conversation_history: list = field(default_factory=list) conversation_history: list = field(default_factory=list)
@@ -61,19 +63,21 @@ async def run_repl(ctx: ReplContext) -> None:
""" """
Enter the interactive REPL. Blocks until the user exits. Enter the interactive REPL. Blocks until the user exits.
""" """
# Build initial scan context string (injected into every AI prompt) # Build scan context string once — injected into every AI prompt
scan_ctx_str = _build_scan_context(ctx) scan_ctx_str = _build_scan_context(ctx)
console_out.print() console.print()
console_out.print("[bold cyan]💬 Ask a follow-up[/bold cyan] [dim](or press Ctrl+C / type /exit to quit)[/dim]") console.rule("[bold cyan] SecureLens AI Chat [/bold cyan]", style="cyan")
console_out.print("[dim]Type /help for available commands[/dim]") console.print(
console_out.print() "[dim]Ask anything about the scan results. "
"Type [bold]/help[/bold] for commands, [bold]Ctrl+C[/bold] to exit.[/dim]\n"
)
while True: while True:
try: try:
user_input = _prompt_user() user_input = Prompt.ask("[bold cyan]>[/bold cyan]")
except (KeyboardInterrupt, EOFError): except (KeyboardInterrupt, EOFError):
console_out.print("\n[dim]Goodbye![/dim]\n") console.print("\n[dim]Goodbye![/dim]\n")
break break
user_input = user_input.strip() user_input = user_input.strip()
@@ -89,44 +93,52 @@ async def run_repl(ctx: ReplContext) -> None:
# ── AI response ───────────────────────────────────────────────────── # ── AI response ─────────────────────────────────────────────────────
if not ctx.api_key: if not ctx.api_key:
print_error("No API key configured. Run `securelens configure` to set one.") console.print(
"\n [bold red]✗ No API key configured.[/bold red] "
"Run [cyan]securelens configure[/cyan] to set one.\n"
)
continue continue
prompt = chat_prompt(ctx.target, scan_ctx_str, user_input) # Show a thinking indicator
console_out.print("[dim] Thinking...[/dim]") with console.status("[dim]Thinking...[/dim]", spinner="dots"):
response = await call_ai( prompt = chat_prompt(ctx.target, scan_ctx_str, user_input)
prompt=prompt, response = await call_ai(
api_key=ctx.api_key, prompt=prompt,
model=ctx.model, api_key=ctx.api_key,
temperature=0.5, model=ctx.model,
conversation_history=ctx.conversation_history, temperature=0.5,
) conversation_history=ctx.conversation_history,
)
if response: if response:
# Save to history for multi-turn context # Append to history for multi-turn context (cap at 20 turns to avoid token bloat)
ctx.conversation_history.append({"role": "user", "content": user_input}) ctx.conversation_history.append({"role": "user", "content": user_input})
ctx.conversation_history.append({"role": "assistant", "content": response}) ctx.conversation_history.append({"role": "assistant", "content": response})
print_ai_response(response) if len(ctx.conversation_history) > 40:
ctx.conversation_history = ctx.conversation_history[-40:]
# Render AI response as Markdown (handles code blocks, bullets, headers)
console.print()
console.print(Panel(
Markdown(response),
border_style="dim cyan",
padding=(0, 1),
))
console.print()
else: else:
print_error("No response from AI. Check your API key and network connection.") console.print(
"\n [bold red]✗ No response from AI.[/bold red] "
"Check your API key and network connection.\n"
)
def _prompt_user() -> str: # ── Scan context builder ──────────────────────────────────────────────────────
"""Read a line from stdin with a styled prompt."""
sys.stdout.write("[dim bold cyan]>[/dim bold cyan] ")
sys.stdout.flush()
# Use input() — rich.prompt not used here to keep it simple
try:
from rich.prompt import Prompt
return Prompt.ask("[bold cyan]>[/bold cyan]")
except Exception:
return input("> ")
def _build_scan_context(ctx: ReplContext) -> str: def _build_scan_context(ctx: ReplContext) -> str:
"""Serialize the scan result into a compact string for the AI context.""" """Serialize the scan result into a compact JSON string for AI context."""
result = ctx.scan_result result = ctx.scan_result
if ctx.target_type == "code":
if ctx.target_type in ("code", "github"):
vulns = [ vulns = [
{ {
"file": v.file_path, "file": v.file_path,
@@ -146,6 +158,7 @@ def _build_scan_context(ctx: ReplContext) -> str:
"vulnerabilities": vulns, "vulnerabilities": vulns,
"ai_summary": result.ai_summary, "ai_summary": result.ai_summary,
}, indent=2) }, indent=2)
else: # web else: # web
issues = [ issues = [
{"layer": i.layer, "severity": i.severity, "issue": i.issue, "fix": i.fix} {"layer": i.layer, "severity": i.severity, "issue": i.issue, "fix": i.fix}
@@ -162,57 +175,119 @@ def _build_scan_context(ctx: ReplContext) -> str:
}, indent=2) }, indent=2)
# ── Slash command dispatcher ──────────────────────────────────────────────────
async def _handle_slash_command(cmd: str, ctx: ReplContext) -> bool: async def _handle_slash_command(cmd: str, ctx: ReplContext) -> bool:
""" """
Process a slash command. Returns True if the REPL should exit. Handle a slash command. Returns True if the REPL should exit.
""" """
parts = cmd.split() parts = cmd.strip().split(maxsplit=2)
command = parts[0].lower() command = parts[0].lower()
if command == "/exit": if command == "/exit":
console_out.print("\n[dim]Goodbye![/dim]\n") console.print("\n[dim]Goodbye![/dim]\n")
return True return True
elif command == "/help": elif command == "/help":
console_out.print(HELP_TEXT) console.print(HELP_TEXT)
elif command == "/clear": elif command == "/clear":
console_out.clear() console.clear()
elif command == "/files": elif command == "/files":
result = ctx.scan_result result = ctx.scan_result
if ctx.target_type == "code" and hasattr(result, "files_triaged"): if ctx.target_type in ("code", "github") and hasattr(result, "files_triaged"):
console_out.print("\n[bold]Files analyzed:[/bold]") if result.files_triaged:
for f in result.files_triaged: console.print("\n[bold]Files analyzed:[/bold]")
console_out.print(f" [dim]• {f}[/dim]") for f in result.files_triaged:
console_out.print() console.print(f" [dim]• {f}[/dim]")
else:
console.print("\n [dim]No files were analyzed.[/dim]")
console.print()
else: else:
print_info("File list not available for web scans.") console.print("\n [dim]File list not available for web scans.[/dim]\n")
elif command == "/score": elif command == "/score":
r = ctx.scan_result r = ctx.scan_result
score = r.score from securelens.output import GRADE_COLOR
grade = r.grade grade_color = GRADE_COLOR.get(r.grade, "white")
console_out.print(f"\n [bold]Score:[/bold] {score}/100 [bold]Grade:[/bold] {grade}\n") console.print(
f"\n Score: [{grade_color}]{r.score}/100 Grade: {r.grade}[/{grade_color}]\n"
)
elif command == "/issues":
severity_filter = parts[1].strip().lower() if len(parts) > 1 else None
_print_issues_summary(ctx, severity_filter)
elif command == "/model": elif command == "/model":
if len(parts) < 2: if len(parts) < 2:
print_info(f"Current model: {ctx.model}") console.print(f"\n [dim]Current model: {ctx.model}[/dim]")
print_info("Usage: /model <model-name> e.g. /model gpt-4o-mini") console.print(" [dim]Usage: /model <model-name> e.g. /model gpt-4o-mini[/dim]\n")
else: else:
ctx.model = parts[1] ctx.model = parts[1]
print_success(f"Model switched to: {ctx.model}") console.print(f"\n [bold green]✓ Model switched to: {ctx.model}[/bold green]\n")
elif command == "/export": elif command == "/export":
fmt = parts[1].lower() if len(parts) > 1 else "markdown" fmt = parts[1].lower() if len(parts) > 1 else "markdown"
target_type = "code" if ctx.target_type in ("code", "github") else "web"
if fmt == "json": if fmt == "json":
path = save_json(ctx.scan_result, ctx.target_type) path = save_json(ctx.scan_result, target_type)
print_success(f"JSON report saved to: {path}") console.print(f"\n [bold green]✓ JSON report saved:[/bold green] [dim]{path}[/dim]\n")
else: else:
path = save_markdown(ctx.scan_result, ctx.target_type) path = save_markdown(ctx.scan_result, target_type)
print_success(f"Markdown report saved to: {path}") console.print(f"\n [bold green]✓ Markdown report saved:[/bold green] [dim]{path}[/dim]\n")
else: else:
print_error(f"Unknown command: {command}. Type /help for available commands.") console.print(
f"\n [bold red]✗ Unknown command: {command}[/bold red] "
"Type [cyan]/help[/cyan] for available commands.\n"
)
return False return False
def _print_issues_summary(ctx: ReplContext, severity_filter: Optional[str] = None) -> None:
"""Print a compact list of all issues, optionally filtered by severity."""
result = ctx.scan_result
from securelens.output import SEVERITY_COLOR
if ctx.target_type in ("code", "github"):
issues = result.vulnerabilities
if not issues:
console.print("\n [bold green]✓ No vulnerabilities found.[/bold green]\n")
return
filtered = issues
if severity_filter:
filtered = [v for v in issues if v.severity.lower() == severity_filter]
if not filtered:
console.print(f"\n [dim]No {severity_filter} issues found.[/dim]\n")
return
console.print(f"\n [bold]{len(filtered)} issue(s):[/bold]")
for i, v in enumerate(filtered, 1):
color = SEVERITY_COLOR.get(v.severity, "white")
loc = v.file_path
if v.line_number:
loc += f":{v.line_number}"
console.print(f" [{color}][{i}] {v.severity}[/{color}] {v.issue} [dim]{loc}[/dim]")
console.print()
else: # web
issues = result.issues
if not issues:
console.print("\n [bold green]✓ No issues found.[/bold green]\n")
return
filtered = issues
if severity_filter:
filtered = [i for i in issues if i.severity.lower() == severity_filter]
if not filtered:
console.print(f"\n [dim]No {severity_filter} issues found.[/dim]\n")
return
console.print(f"\n [bold]{len(filtered)} issue(s):[/bold]")
for i, issue in enumerate(filtered, 1):
color = SEVERITY_COLOR.get(issue.severity, "white")
console.print(f" [{color}][{i}] {issue.severity}[/{color}] {issue.issue} [dim]{issue.layer}[/dim]")
console.print()