template refactor
This commit is contained in:
@@ -1,157 +1,207 @@
|
||||
# utils/reporting_jinja.py
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from html import escape
|
||||
from ipaddress import ip_address
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Set, Union
|
||||
from typing import Dict, List
|
||||
|
||||
from jinja2 import Environment, FileSystemLoader, select_autoescape
|
||||
|
||||
from utils.models import HostResult, HostReport, GroupedReports # <-- add HostReport, GroupedReports
|
||||
from ipaddress import ip_address
|
||||
from utils.models import HostResult, HostReport, GroupedReports
|
||||
|
||||
|
||||
# ---------------------------
|
||||
# Helper render-time functions
|
||||
# ---------------------------
|
||||
|
||||
def fmt_ports(ports: List[int]) -> str:
|
||||
"""
|
||||
Render a simple, comma-separated list of port numbers, deduped and sorted.
|
||||
"""
|
||||
if not ports:
|
||||
return "none"
|
||||
return ", ".join(str(p) for p in sorted(set(int(x) for x in ports)))
|
||||
unique_ports: List[int] = []
|
||||
seen: set[int] = set()
|
||||
for x in ports:
|
||||
p = int(x)
|
||||
if p not in seen:
|
||||
seen.add(p)
|
||||
unique_ports.append(p)
|
||||
|
||||
unique_ports.sort()
|
||||
return ", ".join(str(p) for p in unique_ports)
|
||||
|
||||
|
||||
def badge(text: str, kind: str = "") -> str:
|
||||
"""
|
||||
Return a class-based badge element (no inline color).
|
||||
'kind' should map to CSS variants defined in your stylesheet (e.g., 'ok', 'issue').
|
||||
Unknown kinds fall back to the base 'badge' styling.
|
||||
|
||||
Examples:
|
||||
badge("OK", "ok") -> <span class="badge ok">OK</span>
|
||||
badge("ISSUE", "issue") -> <span class="badge issue">ISSUE</span>
|
||||
badge("info") -> <span class="badge">info</span>
|
||||
"""
|
||||
text_s = escape(str(text))
|
||||
kind_s = str(kind).strip()
|
||||
if kind_s:
|
||||
return f'<span class="badge {kind_s}">{text_s}</span>'
|
||||
return f'<span class="badge">{text_s}</span>'
|
||||
|
||||
def badge(text: str, bg: str, fg: str = "#ffffff") -> str:
|
||||
return (
|
||||
f'<span style="display:inline-block;padding:2px 6px;border-radius:12px;'
|
||||
f'font-size:12px;line-height:16px;background:{bg};color:{fg};'
|
||||
f'font-family:Segoe UI,Arial,sans-serif">{escape(text)}</span>'
|
||||
)
|
||||
|
||||
def pill_state(state: str) -> str:
|
||||
s = (state or "").lower()
|
||||
"""
|
||||
Map port states to neutral/semantic classes (no inline color).
|
||||
Keeps semantics simple so CSS controls appearance.
|
||||
"""
|
||||
s = (state or "").strip().lower()
|
||||
|
||||
if s == "open":
|
||||
return badge("open", "#16a34a")
|
||||
if s in ("open|filtered",):
|
||||
return badge("open|filtered", "#0ea5e9")
|
||||
if s == "filtered":
|
||||
return badge("filtered", "#f59e0b", "#111111")
|
||||
return badge("open", "ok") # green-ish in CSS
|
||||
if s == "closed":
|
||||
return badge("closed", "#ef4444")
|
||||
return badge(s, "#6b7280")
|
||||
return badge("closed", "issue") # red-ish in CSS
|
||||
|
||||
# For less common states, keep a neutral chip/badge so style stays consistent
|
||||
if s in ("open|filtered", "filtered"):
|
||||
return f'<span class="chip">{escape(s)}</span>'
|
||||
|
||||
# Fallback
|
||||
return f'<span class="chip">{escape(s or "-")}</span>'
|
||||
|
||||
|
||||
def pill_proto(proto: str) -> str:
|
||||
return badge((proto or "").lower(), "#334155")
|
||||
"""
|
||||
Render protocol as a neutral 'chip'.
|
||||
"""
|
||||
return f'<span class="chip">{escape((proto or "").lower())}</span>'
|
||||
|
||||
|
||||
def _env(templates_dir: Path) -> Environment:
|
||||
"""
|
||||
Build a Jinja2 environment with our helpers registered as globals.
|
||||
"""
|
||||
env = Environment(
|
||||
loader=FileSystemLoader(str(templates_dir)),
|
||||
autoescape=select_autoescape(["html", "xml"]),
|
||||
autoescape=select_autoescape(["html", "xml", "j2"]),
|
||||
trim_blocks=True,
|
||||
lstrip_blocks=True,
|
||||
)
|
||||
env.globals.update(badge=badge, pill_state=pill_state, pill_proto=pill_proto, fmt_ports=fmt_ports)
|
||||
env.globals.update(
|
||||
badge=badge,
|
||||
pill_state=pill_state,
|
||||
pill_proto=pill_proto,
|
||||
fmt_ports=fmt_ports,
|
||||
)
|
||||
return env
|
||||
|
||||
|
||||
# --- NEW: compatibility shim so renderer can accept either shape ---
|
||||
|
||||
ReportsArg = Union[GroupedReports, Dict[str, Dict[str, List[int]]]]
|
||||
|
||||
def _coerce_to_grouped(reports: ReportsArg) -> GroupedReports:
|
||||
"""
|
||||
Accept either:
|
||||
- GroupedReports (new, preferred), or
|
||||
- legacy dict: { ip: {"unexpected_tcp":[], "missing_tcp":[], "unexpected_udp":[], "missing_udp":[] } }
|
||||
and return a GroupedReports instance.
|
||||
"""
|
||||
if isinstance(reports, GroupedReports):
|
||||
return reports
|
||||
|
||||
# Legacy dict -> build HostReport objects
|
||||
by_ip: Dict[str, HostReport] = {}
|
||||
for ip, d in reports.items():
|
||||
hr = HostReport(
|
||||
ip=str(ip),
|
||||
unexpected_tcp=list(d.get("unexpected_tcp", []) or []),
|
||||
missing_tcp=list(d.get("missing_tcp", []) or []),
|
||||
unexpected_udp=list(d.get("unexpected_udp", []) or []),
|
||||
missing_udp=list(d.get("missing_udp", []) or []),
|
||||
)
|
||||
by_ip[ip] = hr
|
||||
|
||||
# Split & sort by IP
|
||||
issues: List[HostReport] = []
|
||||
expected: List[HostReport] = []
|
||||
for hr in by_ip.values():
|
||||
(issues if hr.has_issues() else expected).append(hr)
|
||||
|
||||
def _ip_key(hr: HostReport):
|
||||
try:
|
||||
return ip_address(hr.ip)
|
||||
except ValueError:
|
||||
return hr.ip
|
||||
|
||||
issues.sort(key=_ip_key)
|
||||
expected.sort(key=_ip_key)
|
||||
|
||||
return GroupedReports(issues=issues, expected=expected, by_ip=by_ip)
|
||||
|
||||
# ---------------------------
|
||||
# Rendering entry points
|
||||
# ---------------------------
|
||||
|
||||
def render_html_report_jinja(
|
||||
reports: ReportsArg,
|
||||
grouped: GroupedReports,
|
||||
host_results: List[HostResult],
|
||||
templates_dir: Path,
|
||||
template_name: str = "report.html.j2",
|
||||
*,
|
||||
template_name: str = "report_body.html.j2",
|
||||
title: str = "Port Compliance Report",
|
||||
only_issues: bool = False,
|
||||
dark_mode: bool = True,
|
||||
) -> str:
|
||||
"""
|
||||
Render the HTML report using the new structure/palette templates.
|
||||
|
||||
Args:
|
||||
grouped: GroupedReports with .issues, .expected, .by_ip already populated.
|
||||
host_results: Raw scan results (per-host) used for the detailed port tables.
|
||||
templates_dir: Path to the Jinja templates root (contains base_report.html.j2, palettes/, etc.).
|
||||
template_name: Structure template to render (defaults to 'report_body.html.j2').
|
||||
title: Report title.
|
||||
only_issues: If True, suppresses the 'expected/OK' section.
|
||||
dark_mode: Controls which palette is included by base_report.html.j2.
|
||||
|
||||
Returns:
|
||||
Rendered HTML as a string.
|
||||
"""
|
||||
env = _env(templates_dir)
|
||||
template = env.get_template(template_name)
|
||||
|
||||
grouped = _coerce_to_grouped(reports)
|
||||
|
||||
total_hosts = len(grouped.by_ip)
|
||||
ok_hosts = len(grouped.expected)
|
||||
hosts_with_issues = [hr.ip for hr in grouped.issues]
|
||||
|
||||
# Build a mapping of IP -> HostResult and sort port rows for stable output
|
||||
by_ip_results = {hr.address: hr for hr in host_results}
|
||||
# Build mapping of IP -> HostResult and sort port rows for stable output
|
||||
by_ip_results: Dict[str, HostResult] = {}
|
||||
for hr in host_results:
|
||||
by_ip_results[hr.address] = hr
|
||||
|
||||
for hr in by_ip_results.values():
|
||||
if hr and hr.ports:
|
||||
hr.ports.sort(key=lambda p: (p.protocol, p.port))
|
||||
# Sort by (protocol, then port, then service)
|
||||
hr.ports.sort(key=lambda p: (str(p.protocol).lower(), int(p.port), str(p.service or "")))
|
||||
|
||||
html = template.render(
|
||||
# Base metadata
|
||||
title=title,
|
||||
generated=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||||
|
||||
# Palette switch (consumed by base_report.html.j2)
|
||||
dark_mode=dark_mode,
|
||||
|
||||
# Summary bar
|
||||
total_hosts=total_hosts,
|
||||
ok_hosts=ok_hosts,
|
||||
hosts_with_issues=hosts_with_issues,
|
||||
|
||||
# New grouped context for simple “issues first, expected later” loops
|
||||
issues=grouped.issues, # list[HostReport]
|
||||
expected=grouped.expected, # list[HostReport]
|
||||
by_ip=grouped.by_ip, # dict[str, HostReport]
|
||||
# Sections
|
||||
issues=grouped.issues, # list[HostReport]
|
||||
expected=grouped.expected, # list[HostReport]
|
||||
by_ip=grouped.by_ip, # dict[str, HostReport]
|
||||
|
||||
# Legacy context (kept for compatibility if your template still uses it)
|
||||
reports=reports,
|
||||
|
||||
# Host scan details
|
||||
# Host scan details (used for port tables)
|
||||
host_results_by_ip=by_ip_results,
|
||||
|
||||
# Existing behavior switch
|
||||
# Behavior flag
|
||||
only_issues=only_issues,
|
||||
)
|
||||
return html
|
||||
|
||||
|
||||
def write_html_report_jinja(
|
||||
reports: ReportsArg,
|
||||
grouped: GroupedReports,
|
||||
host_results: List[HostResult],
|
||||
out_path: Path,
|
||||
*,
|
||||
templates_dir: Path = Path("templates"),
|
||||
template_name: str = "report.html.j2",
|
||||
template_name: str = "report_body.html.j2",
|
||||
title: str = "Port Compliance Report",
|
||||
only_issues: bool = False,
|
||||
dark_mode: bool = True,
|
||||
) -> Path:
|
||||
"""
|
||||
Render and write the HTML report to disk.
|
||||
|
||||
Notes:
|
||||
- Assumes your templates directory contains:
|
||||
base_report.html.j2
|
||||
report_body.html.j2
|
||||
palettes/light.css.j2
|
||||
palettes/dark.css.j2
|
||||
- 'dark_mode' toggles which palette is inlined by the base template.
|
||||
"""
|
||||
html = render_html_report_jinja(
|
||||
reports, host_results, templates_dir, template_name, title, only_issues
|
||||
grouped=grouped,
|
||||
host_results=host_results,
|
||||
templates_dir=templates_dir,
|
||||
template_name=template_name,
|
||||
title=title,
|
||||
only_issues=only_issues,
|
||||
dark_mode=dark_mode,
|
||||
)
|
||||
out_path.write_text(html, encoding="utf-8")
|
||||
return out_path
|
||||
|
||||
Reference in New Issue
Block a user