From 1e8526f4425b70d99d049eacdd0eb427e9de7888 Mon Sep 17 00:00:00 2001 From: Alexandre Pires Date: Sun, 21 Sep 2025 23:55:53 +0200 Subject: [PATCH] Enhance Trivy integration: update email dispatch logic, add LLM prompt generation, and improve security context in deployments --- ansible/host_vars/dev-01.lab.alexpires.me.yml | 1 + terraform/apps/prod-01/config.tf | 13 +- .../apps/nginx-reverse-proxy/config.tf | 19 +- .../modules/apps/nginx-reverse-proxy/main.tf | 50 +-- .../apps/nginx-reverse-proxy/variables.tf | 5 +- terraform/modules/apps/ssh_locker_web/main.tf | 1 + terraform/modules/utils/trivy/main.tf | 44 +- .../utils/trivy/resources/send_report.py | 375 ++++++++++++++++-- .../utils/trivy/resources/send_report_test.py | 16 + 9 files changed, 441 insertions(+), 83 deletions(-) create mode 100644 terraform/modules/utils/trivy/resources/send_report_test.py diff --git a/ansible/host_vars/dev-01.lab.alexpires.me.yml b/ansible/host_vars/dev-01.lab.alexpires.me.yml index dac85ad..ce2bb59 100644 --- a/ansible/host_vars/dev-01.lab.alexpires.me.yml +++ b/ansible/host_vars/dev-01.lab.alexpires.me.yml @@ -40,6 +40,7 @@ fw_allowed_ports: - { rule: "allow", port: "445", proto: "tcp", from: "10.19.4.0/24" } # Samba (local network only) - { rule: "allow", port: "11434", proto: "tcp", from: "10.19.4.0/24" } # Ollama (local network only) - { rule: "allow", port: "11434", proto: "tcp", from: "10.5.5.5/32" } # Ollama (Laptop) + - { rule: "allow", port: "11434", proto: "tcp", from: "10.5.5.2/32" } # Ollama (Contabo VPN) - { rule: "allow", port: "443", proto: "tcp", from: "10.5.5.2/32" } # HTTPS (Contabo VPN) - { rule: "allow", port: "53", proto: "udp", from: "10.19.4.0/24" } # DNS (local network) - { rule: "allow", port: "53", proto: "udp", from: "10.5.5.2/32" } # DNS (Contabo VPN) diff --git a/terraform/apps/prod-01/config.tf b/terraform/apps/prod-01/config.tf index d583f59..0be2a72 100644 --- a/terraform/apps/prod-01/config.tf +++ b/terraform/apps/prod-01/config.tf @@ -166,7 +166,18 @@ locals { }, ] }, - + "ollama" = { + http_port = 11434 + ingress = false + resolver = "10.19.4.136" + upstream = "http://ollama.lab.alexpires.me:11434" + service_http_port = 11434 + locations = [ + { + path = "/" + }, + ] + } } backup_bucket_name = data.terraform_remote_state.scaleway.outputs.backup_bucket_name diff --git a/terraform/modules/apps/nginx-reverse-proxy/config.tf b/terraform/modules/apps/nginx-reverse-proxy/config.tf index 3bfb34c..efbbcb0 100644 --- a/terraform/modules/apps/nginx-reverse-proxy/config.tf +++ b/terraform/modules/apps/nginx-reverse-proxy/config.tf @@ -24,14 +24,17 @@ locals { services_config = { for k, v in var.services : "${k}.conf" => coalesce(v.config, templatefile("${path.module}/resources/nginx/conf.d/default.conf.tpl", { - errors = local.error_codes - http_port = v.http_port - https_port = v.https_port - upstream = v.upstream - resolver = v.resolver - tls = v.tls - fqdn = join(" ", v.fqdn) - auth = length([for l in v.locations : true if l.private]) > 0 + errors = local.error_codes + http_port = v.http_port + https_port = v.https_port + upstream = v.upstream + resolver = v.resolver + ingress = v.ingress + service_http_port = v.service_http_port + service_https_port = v.service_https_port + tls = v.tls + fqdn = length(v.fqdn) > 0 ? join(" ", v.fqdn) : k + auth = length([for l in v.locations : true if l.private]) > 0 locations = [for l in v.locations : { path = l.path headers = merge(l.headers, coalesce(v.default_headers, local.default_headers)) diff --git a/terraform/modules/apps/nginx-reverse-proxy/main.tf b/terraform/modules/apps/nginx-reverse-proxy/main.tf index 6417cc0..20e12bc 100644 --- a/terraform/modules/apps/nginx-reverse-proxy/main.tf +++ b/terraform/modules/apps/nginx-reverse-proxy/main.tf @@ -22,23 +22,23 @@ resource "kubernetes_service_account" "reverse-proxy_service_account" { automount_service_account_token = false } -resource "kubernetes_service" "reverse-proxy" { +resource "kubernetes_service" "services" { for_each = var.services metadata { - name = format("reverse-proxy-%s", replace(each.key, ".", "-")) + name = format("%s", replace(each.key, ".", "-")) namespace = kubernetes_namespace.reverse_proxy.metadata[0].name } spec { port { name = "http" - port = 80 + port = each.value.service_http_port target_port = substr(format("http-%s", replace(each.key, ".", "-")), 0, 15) } dynamic "port" { for_each = each.value.tls ? [1] : [] content { name = "https" - port = 443 + port = each.value.service_https_port target_port = substr(format("https-%s", replace(each.key, ".", "-")), 0, 15) } } @@ -137,10 +137,10 @@ resource "kubernetes_config_map" "auth_script" { } } -resource "kubernetes_ingress_v1" "reverse-proxy_ingress" { - for_each = var.services +resource "kubernetes_ingress_v1" "ingress" { + for_each = { for k, v in var.services : k => v if lookup(v, "ingress", true) && length(v.fqdn) > 0 } metadata { - name = format("reverse-proxy-%s", replace(each.key, ".", "-")) + name = format("%s", replace(each.key, ".", "-")) namespace = kubernetes_namespace.reverse_proxy.metadata[0].name annotations = merge({ @@ -149,7 +149,7 @@ resource "kubernetes_ingress_v1" "reverse-proxy_ingress" { }, each.value.custom_ingress_annotations, each.value.tls ? { - "nginx.org/ssl-services" = format("reverse-proxy-%s", replace(each.key, ".", "-")) + "nginx.org/ssl-services" = format("%s", replace(each.key, ".", "-")) "nginx.org/ssl-redirect" = "true" "nginx.org/redirect-to-https" = "true" "nginx.ingress.kubernetes.io/backend-protocol" = "HTTPS" @@ -169,7 +169,7 @@ resource "kubernetes_ingress_v1" "reverse-proxy_ingress" { path { backend { service { - name = format("reverse-proxy-%s", replace(each.key, ".", "-")) + name = format("%s", replace(each.key, ".", "-")) port { name = each.value.tls ? "https" : "http" } @@ -535,35 +535,3 @@ resource "kubernetes_deployment" "reverse_proxy" { ignore_changes = [spec[0].template[0].metadata[0].annotations["kubectl.kubernetes.io/restartedAt"]] } } - - -# resource "kubernetes_service" "reverse_proxy_streams" { -# count = length(var.streams) > 0 ? 1 : 0 -# metadata { -# name = "stream-ports" -# namespace = kubernetes_namespace.reverse_proxy.metadata[0].name -# } - -# spec { -# selector = { -# app = "reverse-proxy" -# } - -# dynamic "port" { -# for_each = var.streams -# content { -# name = format("%s-%s-%d", port.key, port.value.protocol, port.value.port) -# port = port.value.port -# target_port = port.value.port -# protocol = upper(port.value.protocol) -# } -# } -# type = "LoadBalancer" -# } - -# lifecycle { -# ignore_changes = [ -# metadata[0].annotations -# ] -# } -# } diff --git a/terraform/modules/apps/nginx-reverse-proxy/variables.tf b/terraform/modules/apps/nginx-reverse-proxy/variables.tf index 76fc960..586add7 100644 --- a/terraform/modules/apps/nginx-reverse-proxy/variables.tf +++ b/terraform/modules/apps/nginx-reverse-proxy/variables.tf @@ -29,7 +29,7 @@ variable "issuer_name" { variable "services" { description = "A list of services to be deployed" type = map(object({ - fqdn = list(string) + fqdn = optional(list(string), []) custom_ingress_annotations = optional(map(string), {}) annotations = optional(map(string), {}) upstream = optional(string, "") @@ -38,6 +38,9 @@ variable "services" { http_port = optional(number, 80) https_port = optional(number, 443) tls = optional(bool, false) + ingress = optional(bool, true) + service_http_port = optional(number, 80) + service_https_port = optional(number, 443) locations = optional(list(object({ path = string, headers = optional(map(string), {}), diff --git a/terraform/modules/apps/ssh_locker_web/main.tf b/terraform/modules/apps/ssh_locker_web/main.tf index a378a5b..409309c 100644 --- a/terraform/modules/apps/ssh_locker_web/main.tf +++ b/terraform/modules/apps/ssh_locker_web/main.tf @@ -85,6 +85,7 @@ resource "kubernetes_deployment" "ssh_locker_web" { security_context { allow_privilege_escalation = false + read_only_root_filesystem = true } port { diff --git a/terraform/modules/utils/trivy/main.tf b/terraform/modules/utils/trivy/main.tf index 49cdf43..80b0e4d 100644 --- a/terraform/modules/utils/trivy/main.tf +++ b/terraform/modules/utils/trivy/main.tf @@ -99,7 +99,24 @@ resource "kubernetes_cron_job_v1" "trivy" { 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.txt", "--disable-node-collector"] + 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" + ] + security_context { + run_as_non_root = true + run_as_user = 1000 + run_as_group = 1000 + allow_privilege_escalation = false + read_only_root_filesystem = true + } volume_mount { name = "trivy-reports" mount_path = "/reports" @@ -108,6 +125,10 @@ resource "kubernetes_cron_job_v1" "trivy" { name = "trivy-cache" mount_path = "/cache" } + volume_mount { + name = "tmp" + mount_path = "/tmp" + } } container { name = "send-report" @@ -115,7 +136,7 @@ resource "kubernetes_cron_job_v1" "trivy" { args = ["/scripts/send_report.py"] env { name = "REPORT_PATH" - value = "/reports/trivy.txt" + value = "/reports/trivy.json" } env { name = "SMTP_HOST" @@ -151,6 +172,13 @@ resource "kubernetes_cron_job_v1" "trivy" { name = "EMAIL_TO" value = var.to_email } + security_context { + run_as_non_root = true + run_as_user = 1000 + run_as_group = 1000 + allow_privilege_escalation = false + read_only_root_filesystem = true + } volume_mount { name = "trivy-reports" mount_path = "/reports" @@ -159,6 +187,10 @@ resource "kubernetes_cron_job_v1" "trivy" { name = "trivy-scripts" mount_path = "/scripts" } + volume_mount { + name = "tmp" + mount_path = "/tmp" + } } volume { name = "trivy-reports" @@ -166,9 +198,7 @@ resource "kubernetes_cron_job_v1" "trivy" { } volume { name = "trivy-cache" - host_path { - path = local.cache_path - } + empty_dir {} } volume { name = "trivy-scripts" @@ -178,6 +208,10 @@ resource "kubernetes_cron_job_v1" "trivy" { } } + volume { + name = "tmp" + empty_dir {} + } } } } diff --git a/terraform/modules/utils/trivy/resources/send_report.py b/terraform/modules/utils/trivy/resources/send_report.py index 63f6581..b22b770 100644 --- a/terraform/modules/utils/trivy/resources/send_report.py +++ b/terraform/modules/utils/trivy/resources/send_report.py @@ -2,20 +2,29 @@ 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 -REPORT_PATH = os.environ.get("REPORT_PATH", "/tmp/trivy_report.txt") -WAIT_TIMEOUT = int(os.environ.get("WAIT_TIMEOUT", "200")) -SLEEP_INTERVAL = int(os.environ.get("SLEEP_INTERVAL", "5") ) -def dispatch_email(report_path, subject): +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. + """ 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 report_path: - msg.set_content("Trivy scan report attached.") + 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.") # Attach the report file with open(report_path, "rb") as f: @@ -25,38 +34,350 @@ def dispatch_email(report_path, subject): 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") - with smtplib.SMTP(smtp_host, smtp_port) as server: + # Use context manager but create appropriate SMTP class depending on port + try: if smtp_port == 465: server = smtplib.SMTP_SSL(smtp_host, smtp_port) else: server = smtplib.SMTP(smtp_host, smtp_port) - if smtp_port != 465: - server.starttls() + 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): 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.") + 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 -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_PATH): - if time.time() - start > WAIT_TIMEOUT: - print("Timeout waiting for report to be generated.") - dispatch_email(None, "Trivy Scan Report - Timeout") + +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 + +if __name__ == "__main__": + + 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) - print("Still waiting for report to be generated...") - time.sleep(SLEEP_INTERVAL) -print("Report generated, proceeding to send email...") -time.sleep(5) # Wait for a bit to ensure the file is fully written + # 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("Starting to wait for the report to be generated...") + start = time.time() + while not os.path.exists(report_file): + if time.time() - start > wait_timeout: + print("Timeout waiting for report to be generated.") + dispatch_email(None, "Trivy Scan Report - Timeout") + exit(1) + print("Still waiting for report to be generated...") + time.sleep(sleep_interval) + + print("Report generated, proceeding to process and 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 new file mode 100644 index 0000000..044be99 --- /dev/null +++ b/terraform/modules/utils/trivy/resources/send_report_test.py @@ -0,0 +1,16 @@ +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))