diff --git a/terraform/modules/utils/trivy/config.tf b/terraform/modules/utils/trivy/config.tf index 0423411..5c75ffc 100644 --- a/terraform/modules/utils/trivy/config.tf +++ b/terraform/modules/utils/trivy/config.tf @@ -3,5 +3,6 @@ locals { cronjob_schedule = var.cron_schedule python_version = "3.12-alpine3.21" - send_report_script = file("${path.module}/resources/send_report.py") + send_report_script = file("${path.module}/resources/send_report.py") + generate_report_script = file("${path.module}/resources/generate_report.sh") } diff --git a/terraform/modules/utils/trivy/main.tf b/terraform/modules/utils/trivy/main.tf index 80b0e4d..8b7f1d1 100644 --- a/terraform/modules/utils/trivy/main.tf +++ b/terraform/modules/utils/trivy/main.tf @@ -56,13 +56,14 @@ resource "kubernetes_secret" "email_credentials" { } } -resource "kubernetes_config_map" "send_report" { +resource "kubernetes_config_map" "scripts" { metadata { - name = "trivy-send-report" + name = "trivy-scripts" namespace = kubernetes_namespace.trivy.metadata[0].name } data = { - "send_report.py" = local.send_report_script + "send_report.py" = local.send_report_script + "generate_report.sh" = local.generate_report_script } } @@ -97,19 +98,9 @@ resource "kubernetes_cron_job_v1" "trivy" { service_account_name = kubernetes_service_account.trivy.metadata[0].name automount_service_account_token = true container { - image = "aquasec/trivy:${var.tag}" - name = "trivy" - args = [ - "k8s", - "--report=${var.report_type}", - "--exit-code=0", - "--severity=${var.levels}", - "--cache-dir=/cache", - "-o", "/reports/trivy.json", - "--format", "json", - "--disable-node-collector", - "--parallel", "1" - ] + image = "aquasec/trivy:${var.tag}" + name = "trivy" + command = ["/bin/sh", "-c", "/scripts/generate_report.sh"] security_context { run_as_non_root = true run_as_user = 1000 @@ -117,13 +108,17 @@ resource "kubernetes_cron_job_v1" "trivy" { allow_privilege_escalation = false read_only_root_filesystem = true } - volume_mount { - name = "trivy-reports" - mount_path = "/reports" + env { + name = "REPORT_TYPE" + value = var.report_type + } + env { + name = "LEVELS" + value = var.levels } volume_mount { - name = "trivy-cache" - mount_path = "/cache" + name = "trivy-scripts" + mount_path = "/scripts" } volume_mount { name = "tmp" @@ -134,10 +129,6 @@ resource "kubernetes_cron_job_v1" "trivy" { name = "send-report" image = "python:${local.python_version}" args = ["/scripts/send_report.py"] - env { - name = "REPORT_PATH" - value = "/reports/trivy.json" - } env { name = "SMTP_HOST" value = var.smtp_host @@ -179,10 +170,6 @@ resource "kubernetes_cron_job_v1" "trivy" { allow_privilege_escalation = false read_only_root_filesystem = true } - volume_mount { - name = "trivy-reports" - mount_path = "/reports" - } volume_mount { name = "trivy-scripts" mount_path = "/scripts" @@ -192,18 +179,10 @@ resource "kubernetes_cron_job_v1" "trivy" { mount_path = "/tmp" } } - volume { - name = "trivy-reports" - empty_dir {} - } - volume { - name = "trivy-cache" - empty_dir {} - } volume { name = "trivy-scripts" config_map { - name = kubernetes_config_map.send_report.metadata[0].name + name = kubernetes_config_map.scripts.metadata[0].name default_mode = "0755" } diff --git a/terraform/modules/utils/trivy/resources/generate_report.sh b/terraform/modules/utils/trivy/resources/generate_report.sh new file mode 100644 index 0000000..1ceb0f7 --- /dev/null +++ b/terraform/modules/utils/trivy/resources/generate_report.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env sh +set -euo pipefail +export HOME="/tmp" +TRIVY_CACHE_DIR="/tmp/cache" +mkdir -p /tmp/.trivy +mkdir -p ${TRIVY_CACHE_DIR} +trivy plugin install scan2html +trivy k8s --report=${REPORT_TYPE:-all} \ + --exit-code=0 \ + --severity=${LEVELS:-CRITICAL,HIGH,MEDIUM} \ + -o tmp/trivy.json \ + --format json \ + --disable-node-collector \ + --parallel ${PARALLEL:-5} +trivy scan2html generate --scan2html-flags --output tmp/report.html --from tmp/trivy.json \ No newline at end of file diff --git a/terraform/modules/utils/trivy/resources/send_report.py b/terraform/modules/utils/trivy/resources/send_report.py index b22b770..17da03f 100644 --- a/terraform/modules/utils/trivy/resources/send_report.py +++ b/terraform/modules/utils/trivy/resources/send_report.py @@ -2,29 +2,16 @@ import os import time import smtplib -import json -import uuid -import ssl -import urllib.request -import urllib.error from email.message import EmailMessage -from typing import List, Dict, Any, Optional - -def dispatch_email(report_path: str, subject: str, extra_attachments: List[str] = None, body: Optional[str] = None): - """Send an email with the Trivy report and optional extra attachments (like prompts JSONL). - - Attachments are optional and will be ignored if files are missing. - """ +def dispatch_email(report_path, subject): msg = EmailMessage() msg["Subject"] = subject msg["From"] = os.environ.get("EMAIL_FROM", "noreply@example.com") msg["To"] = os.environ.get("EMAIL_TO", "admin@example.com") - if body: - msg.set_content(body) - elif report_path and os.path.exists(report_path): - msg.set_content("Trivy scan report attached. Prompts file may also be attached if configured.") + if report_path: + msg.set_content("Trivy scan report attached.") # Attach the report file with open(report_path, "rb") as f: @@ -34,308 +21,35 @@ def dispatch_email(report_path: str, subject: str, extra_attachments: List[str] else: msg.set_content("Trivy scan report not found. Check logs for details.") - # Attach any extra files (prompts etc) - if extra_attachments: - for path in extra_attachments: - try: - if not path or not os.path.exists(path): - continue - with open(path, "rb") as f: - data = f.read() - name = os.path.basename(path) - # Use generic octet-stream for unknown types - msg.add_attachment(data, maintype="application", subtype="octet-stream", filename=name) - except Exception as e: - print(f"Warning: failed to attach {path}: {e}") - smtp_host = os.environ.get("SMTP_HOST", "smtp.example.com") smtp_port = int(os.environ.get("SMTP_PORT", "587")) smtp_user = os.environ.get("SMTP_USER", "user") smtp_pass = os.environ.get("SMTP_PASS", "pass") - # Use context manager but create appropriate SMTP class depending on port - try: + with smtplib.SMTP(smtp_host, smtp_port) as server: if smtp_port == 465: server = smtplib.SMTP_SSL(smtp_host, smtp_port) else: server = smtplib.SMTP(smtp_host, smtp_port) - with server: - if smtp_port != 465: - server.starttls() - try: - server.login(smtp_user, smtp_pass) - except smtplib.SMTPAuthenticationError: - print("SMTP Authentication Error: Check your SMTP credentials.") - return - server.send_message(msg) - print("Report sent.") - except Exception as e: - print(f"Failed to send email: {e}") - - -def load_trivy_report(path: str) -> Dict[str, Any]: - with open(path, "r") as f: - return json.load(f) - - -def make_prompt_for_findings(resource_id: str, findings: List[Dict[str, Any]]) -> str: - """Create an LLM-ready prompt for a list of findings for a single resource. - - The prompt requests a concise summary, prioritized remediation steps, and suggested changes. - """ - lines = [f"Trivy scan findings for resource: {resource_id}", "\nPlease analyze the following findings:"] - for i, fnd in enumerate(findings, 1): - title = fnd.get("Title") or fnd.get("PkgName") or fnd.get("VulnerabilityID") or fnd.get("Type") - sev = fnd.get("Severity") or fnd.get("SeveritySource") or "UNKNOWN" - desc = fnd.get("Description") or fnd.get("Message") or "" - extra = [] - if "VulnerabilityID" in fnd: - extra.append(f"id={fnd.get('VulnerabilityID')}") - if "InstalledVersion" in fnd: - extra.append(f"installed={fnd.get('InstalledVersion')}") - if "FixedVersion" in fnd: - extra.append(f"fixed={fnd.get('FixedVersion')}") - lines.append(f"{i}. {title} (severity={sev}) {' '.join(extra)}") - if desc: - # keep description short - lines.append(" " + (desc.strip().split('\n')[0]) ) - - lines.append("\nPlease output a JSON object with keys: summary, prioritized_remediation (list), suggested_changes (concrete steps or code/manifest snippets), references (list). Keep answers concise.") - return "\n".join(lines) - - -def generate_llm_prompts(report_path: str, output_path: str, individual_severities: List[str], batch_size: int = 10) -> List[str]: - """Parse Trivy JSON and generate a JSONL file with prompts. - - Returns list of generated prompt file paths (currently only one file). - """ - report = load_trivy_report(report_path) - resources = report.get("Resources", []) if isinstance(report, dict) else [] - - prompts = [] - - # Collect lower severity findings grouped per resource - grouped_by_resource: Dict[str, List[Dict[str, Any]]] = {} - - for res in resources: - kind = res.get("Kind", "") - name = res.get("Name", "") - resource_id = f"{kind}/{name}" if name else kind - - for result in res.get("Results", []) or []: - # Trivy stores vulnerabilities under 'Vulnerabilities' and misconfigs under 'Misconfigurations' etc. - for k in ("Vulnerabilities", "Misconfigurations", "Secrets", "MisconfigurationResults"): - for item in result.get(k, []) if isinstance(result.get(k, []), list) else []: - sev = (item.get("Severity") or item.get("SeveritySource") or "UNKNOWN").upper() - # If severity in the individual list, make a dedicated prompt - if sev in individual_severities: - prompt_text = make_prompt_for_findings(resource_id, [item]) - prompts.append({ - "id": str(uuid.uuid4()), - "resource": resource_id, - "severity": sev, - "prompt": prompt_text, - "metadata": item, - }) - else: - grouped_by_resource.setdefault(resource_id, []).append(item) - - # Now batch grouped findings per resource into chunks of batch_size - for resource_id, items in grouped_by_resource.items(): - for i in range(0, len(items), batch_size): - chunk = items[i : i + batch_size] - # pick highest severity in chunk for metadata - highest = "UNKNOWN" - for it in chunk: - s = (it.get("Severity") or it.get("SeveritySource") or "").upper() - if s == "CRITICAL": - highest = "CRITICAL" - break - if s == "HIGH": - highest = "HIGH" - prompt_text = make_prompt_for_findings(resource_id, chunk) - prompts.append({ - "id": str(uuid.uuid4()), - "resource": resource_id, - "severity": highest, - "prompt": prompt_text, - "metadata": {"count": len(chunk)}, - }) - - # Write JSONL file - try: - with open(output_path, "w") as out: - for p in prompts: - out.write(json.dumps(p) + "\n") - except Exception as e: - print(f"Failed to write prompts file: {e}") - return [] - - return [output_path] - - -def wake_ollama() -> bool: - """Trigger model loading by making a lightweight generate call to Ollama. - - We send a tiny prompt asking the model to respond briefly. If a response (or non-error) is received, we consider the model awake. - """ - wake_timeout = int(os.environ.get("OLLAMA_WAKE_TIMEOUT", "5")) - wake_retries = int(os.environ.get("OLLAMA_WAKE_RETRIES", "6")) - base_url = os.environ.get("OLLAMA_BASE_URL", "http://ollama.reverse-proxy.svc.cluster.local:11434") - last_err = None - for attempt in range(1, wake_retries + 1): + if smtp_port != 465: + server.starttls() try: - req = urllib.request.Request(base_url, method="GET") - with urllib.request.urlopen(req, timeout=wake_timeout) as resp: - body = resp.read().decode(errors="ignore") - if body != "Ollama is running": - print(f"Wake attempt {attempt} unexpected response body: {body}") - else: - print(f"Ollama is running (attempt {attempt})") - return True - except urllib.error.HTTPError as e: - last_err = f"HTTPError {e.code}: {e.reason}" - print(f"Ollama wake HTTPError: {e}") - sleep = min(5 * attempt, 30) - print(f"Retrying wake in {sleep}s...") - time.sleep(sleep) - print(f"Ollama wake failed after {wake_retries} attempts: last_err={last_err}") - return False - - -def call_ollama_generate(prompt: str) -> Dict[str, Any]: - """Call Ollama generate endpoint. Returns parsed JSON response or {'error': str}. - - The function is deliberately conservative about response parsing: if typical keys are present, it extracts the text, otherwise returns the full JSON. - """ - base_url = os.environ.get("OLLAMA_BASE_URL", "http://ollama.reverse-proxy.svc.cluster.local:11434") - model = os.environ.get("OLLAMA_MODEL", "mistral:latest") - timeout = int(os.environ.get("OLLAMA_TIMEOUT", "30")) - retries = int(os.environ.get("OLLAMA_RETRIES", "2")) - - url = base_url.rstrip("/") + "/api/generate" - payload = {"model": model, "prompt": prompt, "stream": False, "raw": False} - data = json.dumps(payload).encode() - last_err = None - for attempt in range(1, retries + 1): - try: - req = urllib.request.Request(url, data=data, method="POST") - with urllib.request.urlopen(req, timeout=timeout) as resp: - body = resp.read().decode(errors="ignore") - try: - j = json.loads(body) - except Exception: - return {"error": "Failed to parse JSON response:"} - - try: - llm_response = j.get("response", "{}") - llm_json = json.loads(llm_response) - return llm_json - except Exception: - return {"error": "Failed to parse JSON response"} - except urllib.error.HTTPError as e: - last_err = f"HTTPError {e.code}: {e.reason}" - print(f"Ollama call HTTPError: {e}") - except Exception as e: - last_err = str(e) - print(f"Ollama call failed: {e}") - time.sleep(1 + attempt) - return {"error": last_err} - - -def process_prompts_with_ollama(prompts_file: str, max_prompts: Optional[int] = None) -> List[Dict[str, Any]]: - """Read prompts JSONL and call Ollama for each prompt. Returns list of results. - - Each result item contains: id, resource, severity, prompt, response (dict) - """ - # Optionally wake Ollama BEFORE processing the report (we trigger a small generate request to load the model) - max_prompts = os.environ.get("OLLAMA_MAX_PROMPTS", None) - if max_prompts is not None: - max_prompts = int(max_prompts) - - results = [] - if not os.path.exists(prompts_file): - print(f"Prompts file not found: {prompts_file}") - return results - - with open(prompts_file, "r") as f: - for idx, line in enumerate(f): - if max_prompts is not None and idx >= max_prompts: - break - try: - p = json.loads(line) - except Exception as e: - print(f"Failed to parse prompt line {idx}: {e}") - continue - prompt_text = p.get("prompt") or json.dumps(p) - resp = call_ollama_generate(prompt_text) - results.append({ - "id": p.get("id"), - "resource": p.get("resource"), - "severity": p.get("severity"), - "prompt": prompt_text, - "response": resp, - }) - return results - - -def get_email_body_lines(responses: Dict[str, Any]) -> List[str]: - lines = ["Trivy scan report attached."] - if responses: - lines.append("\nLLM Analysis Summary:\n") - for r in responses: - res_id = r.get("resource", "unknown resource") - sev = r.get("severity", "UNKNOWN") - resp = r.get("response", {}) - if "error" in resp: - lines.append(f"- {res_id} (severity={sev}): LLM Error: {resp['error']}") - continue - summary = resp.get("summary", "No summary provided.") - remediation = resp.get("prioritized_remediation", []) - suggestions = resp.get("suggested_changes", []) - references = resp.get("references", []) - - lines.append(f"- {res_id} (severity={sev}):") - lines.append(f" Summary: {summary}") - if remediation: - lines.append(" Prioritized Remediation Steps:") - for step in remediation: - lines.append(f" - {step}") - if suggestions: - lines.append(" Suggested Changes:") - for change in suggestions: - lines.append(f" - {change}") - if references: - lines.append(" References:") - for ref in references: - lines.append(f" - {ref}") - lines.append("") # Blank line between entries - else: - lines.append("\nNo LLM analysis was performed or no responses received.") - return lines + server.login(smtp_user, smtp_pass) + except smtplib.SMTPAuthenticationError: + print("SMTP Authentication Error: Check your SMTP credentials.") + return + server.send_message(msg) + print("Report sent.") if __name__ == "__main__": + report_path = os.environ.get("REPORT_PATH", "/tmp/report.html") + wait_timeout = int(os.environ.get("WAIT_TIMEOUT", "200")) + sleep_interval = int(os.environ.get("SLEEP_INTERVAL", "5") ) - report_file = os.environ.get("REPORT_PATH", None) - wait_timeout = int(os.environ.get("WAIT_TIMEOUT", "200")) - sleep_interval = int(os.environ.get("SLEEP_INTERVAL", "5")) - - if not report_file: - print("REPORT_PATH environment variable is not set.") - exit(1) - - # Do a quick wake; best-effort - if it fails we'll still send report but LLM processing will also be skipped - woke = False - try: - woke = wake_ollama() - except Exception as e: - print(f"Ollama wake failed: {e}") - woke = False - - print("Trivy report path:", report_file) + print("Trivy report path:", report_path) print("Starting to wait for the report to be generated...") start = time.time() - while not os.path.exists(report_file): + while not os.path.exists(report_path): if time.time() - start > wait_timeout: print("Timeout waiting for report to be generated.") dispatch_email(None, "Trivy Scan Report - Timeout") @@ -343,41 +57,7 @@ if __name__ == "__main__": print("Still waiting for report to be generated...") time.sleep(sleep_interval) - print("Report generated, proceeding to process and send email...") + print("Report generated, proceeding to send email...") time.sleep(5) # Wait for a bit to ensure the file is fully written - # Prepare prompts - prompts_output = os.environ.get("PROMPTS_OUTPUT", os.path.join(os.path.dirname(report_file), "trivy_prompts.jsonl")) - indiv = os.environ.get("PROMPT_INDIVIDUAL_SEVERITIES", "CRITICAL,HIGH") - individual_severities = [s.strip().upper() for s in indiv.split(",") if s.strip()] - batch_size = int(os.environ.get("PROMPT_BATCH_SIZE", "10")) - - generated = [] - try: - generated = generate_llm_prompts(report_file, prompts_output, individual_severities, batch_size) - print(f"Generated prompts file(s): {generated}") - except Exception as e: - print(f"Failed to generate prompts: {e}") - - attach_prompts = os.environ.get("ATTACH_PROMPTS", "true").lower() in ("1", "true", "yes") - extras = generated if attach_prompts else [] - - # Process prompts through Ollama if token present and wake succeeded - responses = [] - if woke and generated: - try: - responses = process_prompts_with_ollama(generated[0]) - print(f"Obtained {len(responses)} LLM responses") - except Exception as e: - print(f"Failed to process prompts with Ollama: {e}") - responses = [] - else: - if not woke: - print("Ollama not awake, skipping LLM processing") - - # Aggregate responses into an email body (concise) - - body_text = "\n".join(get_email_body_lines(responses)) - - dispatch_email(report_file, "Trivy Scan Report - " + time.strftime("%Y-%m-%d %H:%M:%S"), extra_attachments=extras, body=body_text) - + dispatch_email(report_path, "Trivy Scan Report - " + time.strftime("%Y-%m-%d %H:%M:%S")) diff --git a/terraform/modules/utils/trivy/resources/send_report_test.py b/terraform/modules/utils/trivy/resources/send_report_test.py deleted file mode 100644 index 044be99..0000000 --- a/terraform/modules/utils/trivy/resources/send_report_test.py +++ /dev/null @@ -1,16 +0,0 @@ -import os -from send_report import generate_llm_prompts, wake_ollama, process_prompts_with_ollama, get_email_body_lines - - -PATH = os.path.dirname(os.path.abspath(__file__)) -os.environ["OLLAMA_MAX_PROMPTS"] = "20" -os.environ["OLLAMA_BASE_URL"] = "http://localhost:11434" -os.environ["OLLAMA_MODEL"] = "mistral:latest" - -prompts = generate_llm_prompts(f"{PATH}/fixtures/trivy.json","/tmp/trivy_prompts.jsonl",["CRITICAL","HIGH"],10) -print(prompts) -wake_response = wake_ollama() -print(wake_response) -responses = process_prompts_with_ollama("/tmp/trivy_prompts.jsonl") -email_body_lines = get_email_body_lines(responses) -print("\n".join(email_body_lines))