Skip to content

Commit be901e3

Browse files
Kim-San04claude
andcommitted
feat: add CVSS v3.1 calculator module
reporting/cvss_calculator.py: full CVSS v3.1 spec implementation with pre-built vectors for common finding types (FTP, SMB, RDP, sensitive files). score_finding() heuristic auto-assigns CVSS from finding title/severity. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent b601a29 commit be901e3

21 files changed

Lines changed: 2149 additions & 187 deletions

core/engine.py

Lines changed: 69 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -1,65 +1,45 @@
11
import logging
2+
import os
23
import time
34
from typing import Optional
45

56
from rich.console import Console
67
from rich.panel import Panel
7-
from rich.progress import (
8-
BarColumn,
9-
Progress,
10-
SpinnerColumn,
11-
TaskProgressColumn,
12-
TextColumn,
13-
TimeElapsedColumn,
14-
)
8+
from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn, TimeElapsedColumn
159
from rich.rule import Rule
1610
from rich.table import Table
1711

1812
from core.session import Session
19-
from core.target import Target
2013

2114
logger = logging.getLogger("pentest-agent.engine")
2215
console = Console()
2316

2417

2518
class Engine:
26-
def __init__(
27-
self,
28-
session: Session,
29-
passive_only: bool = False,
30-
delay: float = 1.0,
31-
):
19+
def __init__(self, session: Session, passive_only: bool = False, delay: float = 1.0,
20+
lang: str = "EN", no_pdf: bool = False, no_html: bool = False):
3221
self.session = session
3322
self.passive_only = passive_only
3423
self.delay = delay
24+
self.lang = lang
25+
self.no_pdf = no_pdf
26+
self.no_html = no_html
3527

3628
def run(self):
3729
target = self.session.target
38-
console.print(
39-
Panel(
40-
f"[bold cyan]Target:[/bold cyan] {target}\n"
41-
f"[bold cyan]Type:[/bold cyan] {target.target_type}\n"
42-
f"[bold cyan]Mode:[/bold cyan] {'Passive only' if self.passive_only else 'Full (passive + active)'}",
43-
title="[bold green]Scan Configuration[/bold green]",
44-
border_style="green",
45-
)
46-
)
30+
console.print(Panel(
31+
f"[bold cyan]Target:[/bold cyan] {target}\n"
32+
f"[bold cyan]Type:[/bold cyan] {target.target_type}\n"
33+
f"[bold cyan]Mode:[/bold cyan] {'Passive only' if self.passive_only else 'Full (passive + active)'}\n"
34+
f"[bold cyan]Lang:[/bold cyan] {self.lang}",
35+
title="[bold green]Scan Configuration[/bold green]", border_style="green",
36+
))
4737

4838
steps = self._build_steps()
49-
50-
with Progress(
51-
SpinnerColumn(),
52-
TextColumn("[progress.description]{task.description}"),
53-
BarColumn(),
54-
TaskProgressColumn(),
55-
TimeElapsedColumn(),
56-
console=console,
57-
transient=False,
58-
) as progress:
59-
overall = progress.add_task(
60-
"[cyan]Overall progress", total=len(steps)
61-
)
62-
39+
with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"),
40+
BarColumn(), TaskProgressColumn(), TimeElapsedColumn(),
41+
console=console, transient=False) as progress:
42+
overall = progress.add_task("[cyan]Overall progress", total=len(steps))
6343
for step_name, step_fn in steps:
6444
progress.update(overall, description=f"[cyan]{step_name}")
6545
logger.info(f"Starting: {step_name}")
@@ -96,64 +76,58 @@ def _build_steps(self):
9676
]
9777
return steps
9878

99-
# ── Passive steps ──────────────────────────────────────────────────────
79+
# ── Passive ────────────────────────────────────────────────────────────
10080

10181
def _step_whois(self):
10282
from modules.recon.passive import PassiveRecon
103-
pr = PassiveRecon(self.session.target, self.delay)
104-
data = pr.whois_lookup()
105-
self.session.update_results("passive_recon", "whois", data)
83+
self.session.update_results("passive_recon", "whois",
84+
PassiveRecon(self.session.target, self.delay).whois_lookup())
10685

10786
def _step_dns(self):
10887
from modules.recon.passive import PassiveRecon
109-
pr = PassiveRecon(self.session.target, self.delay)
110-
data = pr.dns_enumeration()
111-
self.session.update_results("passive_recon", "dns", data)
88+
self.session.update_results("passive_recon", "dns",
89+
PassiveRecon(self.session.target, self.delay).dns_enumeration())
11290

11391
def _step_subdomains(self):
11492
from modules.recon.passive import PassiveRecon
11593
from pathlib import Path
116-
pr = PassiveRecon(self.session.target, self.delay)
11794
wordlist = Path(__file__).parent.parent / "wordlists" / "subdomains.txt"
118-
data = pr.subdomain_enumeration(str(wordlist))
119-
self.session.update_results("passive_recon", "subdomains", data)
95+
self.session.update_results("passive_recon", "subdomains",
96+
PassiveRecon(self.session.target, self.delay).subdomain_enumeration(str(wordlist)))
12097

12198
def _step_shodan(self):
12299
from modules.recon.passive import PassiveRecon
123-
pr = PassiveRecon(self.session.target, self.delay)
124-
data = pr.shodan_lookup()
125-
self.session.update_results("passive_recon", "shodan", data)
100+
self.session.update_results("passive_recon", "shodan",
101+
PassiveRecon(self.session.target, self.delay).shodan_lookup())
126102

127103
def _step_http_tech(self):
128104
from modules.recon.passive import PassiveRecon
129-
pr = PassiveRecon(self.session.target, self.delay)
130-
data = pr.detect_web_technologies()
131-
self.session.update_results("passive_recon", "http_tech", data)
105+
self.session.update_results("passive_recon", "http_tech",
106+
PassiveRecon(self.session.target, self.delay).detect_web_technologies())
132107

133-
# ── Active steps ───────────────────────────────────────────────────────
108+
# ── Active ─────────────────────────────────────────────────────────────
134109

135110
def _step_host_discovery(self):
136111
from modules.scanner.port_scanner import PortScanner
137-
ps = PortScanner()
138-
hosts = ps.host_discovery(self.session.target.get_hosts())
112+
hosts = PortScanner().host_discovery(self.session.target.get_hosts())
139113
self.session.update_results("active_recon", "live_hosts", hosts)
140114

141115
def _step_port_scan(self):
142116
from modules.scanner.port_scanner import PortScanner
143117
ps = PortScanner()
144118
live = self.session.results["active_recon"].get("live_hosts", self.session.target.get_hosts())
145-
results = {}
146-
for host in live:
147-
results[host] = ps.port_scan(host)
148-
self.session.update_results("active_recon", "ports", results)
119+
self.session.update_results("active_recon", "ports",
120+
{host: ps.port_scan(host) for host in live})
149121

150122
def _step_service_detection(self):
151123
from modules.scanner.port_scanner import PortScanner
152124
ps = PortScanner()
153-
ports_data = self.session.results["active_recon"].get("ports", {})
154125
results = {}
155-
for host, port_info in ports_data.items():
156-
open_ports = [p for p, info in port_info.items() if info.get("state") == "open"]
126+
for host, port_info in self.session.results["active_recon"].get("ports", {}).items():
127+
if not isinstance(port_info, dict):
128+
continue
129+
open_ports = [p for p, info in port_info.items()
130+
if isinstance(info, dict) and info.get("state") == "open"]
157131
if open_ports:
158132
results[host] = ps.service_detection(host, open_ports)
159133
self.session.update_results("active_recon", "services", results)
@@ -166,11 +140,14 @@ def _step_web_enum(self):
166140
ws = WebScanner(str(wordlist))
167141
results = {}
168142
for host, svc_info in services.items():
143+
if not isinstance(svc_info, dict) or "error" in svc_info:
144+
continue
169145
for port, info in svc_info.items():
146+
if not isinstance(info, dict):
147+
continue
170148
if port in ("80", "443", "8080", "8443") or info.get("name") in ("http", "https", "ssl/http"):
171149
scheme = "https" if port in ("443", "8443") else "http"
172-
url = f"{scheme}://{host}:{port}"
173-
results[url] = ws.full_scan(url)
150+
results[f"{scheme}://{host}:{port}"] = ws.full_scan(f"{scheme}://{host}:{port}")
174151
self.session.update_results("active_recon", "web_enum", results)
175152

176153
def _step_smb_enum(self):
@@ -179,6 +156,8 @@ def _step_smb_enum(self):
179156
ss = ServiceScanner()
180157
results = {}
181158
for host, svc_info in services.items():
159+
if not isinstance(svc_info, dict) or "error" in svc_info:
160+
continue
182161
if "445" in svc_info or "139" in svc_info:
183162
results[host] = ss.smb_enum(host)
184163
self.session.update_results("active_recon", "smb", results)
@@ -189,6 +168,8 @@ def _step_ldap_enum(self):
189168
ss = ServiceScanner()
190169
results = {}
191170
for host, svc_info in services.items():
171+
if not isinstance(svc_info, dict) or "error" in svc_info:
172+
continue
192173
if "389" in svc_info or "636" in svc_info:
193174
results[host] = ss.ldap_enum(host)
194175
self.session.update_results("active_recon", "ldap", results)
@@ -199,11 +180,7 @@ def _run_ai_analysis(self):
199180
console.print(Rule("[bold magenta]AI Analysis[/bold magenta]"))
200181
try:
201182
from modules.ai.analyzer import AIAnalyzer
202-
analyzer = AIAnalyzer()
203-
analysis = analyzer.analyze(
204-
str(self.session.target),
205-
self.session.results,
206-
)
183+
analysis = AIAnalyzer().analyze(str(self.session.target), self.session.results)
207184
self.session.update_results("ai_analysis", "findings", analysis)
208185
console.print("[green] ✓ AI analysis complete[/green]")
209186
except Exception as exc:
@@ -214,8 +191,7 @@ def _generate_reports(self):
214191
console.print(Rule("[bold blue]Generating Reports[/bold blue]"))
215192
try:
216193
from reporting.generator import ReportGenerator
217-
gen = ReportGenerator(self.session)
218-
gen.generate()
194+
ReportGenerator(self.session, lang=self.lang, no_pdf=self.no_pdf, no_html=self.no_html).generate()
219195
console.print(f"[green] ✓ Reports saved to: {self.session.output_dir}[/green]")
220196
except Exception as exc:
221197
logger.warning(f"Report generation failed: {exc}")
@@ -226,37 +202,38 @@ def _print_summary(self):
226202
console.print(Rule("[bold green]Scan Summary[/bold green]"))
227203
ai = self.session.results.get("ai_analysis", {}).get("findings", {})
228204
score = ai.get("attack_surface_score", "N/A")
229-
summary = ai.get("executive_summary", "No AI summary available.")
230205

231206
table = Table(show_header=True, header_style="bold cyan")
232207
table.add_column("Metric", style="cyan")
233208
table.add_column("Value", style="white")
234209

235210
duration = self.session.results.get("scan_duration_seconds", 0)
236-
passive = self.session.results.get("passive_recon", {})
237211
active = self.session.results.get("active_recon", {})
212+
passive = self.session.results.get("passive_recon", {})
238213

239-
dns_records = sum(
240-
len(v) for v in passive.get("dns", {}).values() if isinstance(v, list)
241-
)
242-
live_hosts = len(active.get("live_hosts", [self.session.target.value]))
214+
live_hosts = active.get("live_hosts", [self.session.target.value])
243215
open_ports = sum(
244-
len([p for p, i in ports.items() if i.get("state") == "open"])
245-
for ports in active.get("ports", {}).values()
216+
1 for ports in active.get("ports", {}).values()
217+
if isinstance(ports, dict)
218+
for p, i in ports.items()
219+
if isinstance(i, dict) and i.get("state") == "open"
246220
)
247-
subdomains = len(passive.get("subdomains", []))
221+
dns_count = sum(len(v) for v in passive.get("dns", {}).values() if isinstance(v, list))
222+
findings = ai.get("critical_findings", [])
223+
crit = sum(1 for f in findings if f.get("severity") == "Critical")
224+
high = sum(1 for f in findings if f.get("severity") == "High")
225+
226+
score_str = f"[bold red]{score}[/bold red]" if isinstance(score, int) and score >= 60 else str(score)
248227

249228
table.add_row("Duration", f"{duration:.1f}s")
250-
table.add_row("Live hosts", str(live_hosts))
229+
table.add_row("Live hosts", str(len(live_hosts)))
251230
table.add_row("Open ports", str(open_ports))
252-
table.add_row("DNS records", str(dns_records))
253-
table.add_row("Subdomains found", str(subdomains))
254-
table.add_row("Attack surface score", f"[bold red]{score}[/bold red]" if isinstance(score, int) and score >= 60 else str(score))
231+
table.add_row("DNS records", str(dns_count))
232+
table.add_row("Subdomains found", str(len(passive.get("subdomains", []))))
233+
table.add_row("Findings", f"{len(findings)} total ({crit} critical, {high} high)")
234+
table.add_row("Attack surface score", score_str)
235+
table.add_row("Reports", str(self.session.output_dir))
255236

256237
console.print(table)
257-
console.print(
258-
Panel(summary, title="[bold]Executive Summary[/bold]", border_style="green")
259-
)
260-
console.print(
261-
f"\n[bold green]Output directory:[/bold green] {self.session.output_dir}"
262-
)
238+
summary = ai.get("executive_summary", "No AI summary available.")
239+
console.print(Panel(summary, title="[bold]Executive Summary[/bold]", border_style="green"))

dashboard/__init__.py

Whitespace-only changes.

0 commit comments

Comments
 (0)