#!/usr/bin/env python3
import os
import time
import smtplib
from email.message import EmailMessage

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):
    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.")

        # Attach the report file
        with open(report_path, "rb") as f:
            file_data = f.read()
            file_name = os.path.basename(report_path)
            msg.add_attachment(file_data, maintype="text", subtype="plain", filename=file_name)
    else:
        msg.set_content("Trivy scan report not found. Check logs for details.")

    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:
        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()
        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.")

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")
        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

dispatch_email(REPORT_PATH, "Trivy Scan Report - " + time.strftime("%Y-%m-%d %H:%M:%S"))
