dmarc_visualizer.py

import os
import xml.etree.ElementTree as ET
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime, timezone
import socket
from functools import lru_cache

@lru_cache(maxsize=1000)
def reverse_dns(ip):
    try:
        hostname = socket.gethostbyaddr(ip)[0]
    except socket.herror:
        return ip  # fallback to IP if no reverse DNS

    # Group certain domains
    if hostname.endswith("google.com"):
        return "Google"
    elif hostname.endswith("outlook.com") or hostname.endswith("hotmail.com") or "microsoft" in hostname:
        return "Microsoft"
    else:
        return hostname


def unix_to_date(ts):
    return datetime.fromtimestamp(int(ts), tz=timezone.utc).date()

def parse_report(file_path):
    data = []
    try:
        tree = ET.parse(file_path)
        root = tree.getroot()

        begin = root.find("report_metadata/date_range/begin").text
        date = unix_to_date(begin)

        for record in root.findall("record"):
            row = record.find("row")
            source_ip = row.findtext("source_ip")
            count = int(row.findtext("count"))
            disposition = row.find("policy_evaluated").findtext("disposition")
            dkim_result = row.find("policy_evaluated").findtext("dkim")
            spf_result = row.find("policy_evaluated").findtext("spf")

            data.append({
                "date": date,
                "source_ip": source_ip,
                "count": count,
                "dkim": dkim_result,
                "spf": spf_result,
                "disposition": disposition
            })

    except Exception as e:
        print(f"Error parsing {file_path}: {e}")

    return data

def load_all_reports(folder):
    all_data = []
    for file in os.listdir(folder):
        if file.endswith(".xml"):
            full_path = os.path.join(folder, file)
            all_data.extend(parse_report(full_path))
    return pd.DataFrame(all_data)

def plot_auth_results(df):
    df["dkim_pass"] = df["dkim"] == "pass"
    df["spf_pass"] = df["spf"] == "pass"
    grouped = df.groupby("date").agg({
        "dkim_pass": "sum",
        "spf_pass": "sum",
        "count": "sum"
    }).reset_index()

    plt.figure(figsize=(10, 5))
    plt.plot(grouped["date"], grouped["dkim_pass"], label="DKIM Pass")
    plt.plot(grouped["date"], grouped["spf_pass"], label="SPF Pass")
    plt.plot(grouped["date"], grouped["count"], label="Total Emails", linestyle="--")
    plt.title("DMARC Authentication Trends")
    plt.xlabel("Date")
    plt.ylabel("Count")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.show()

def plot_source_ip_volume(df):
    df["hostname"] = df["source_ip"].apply(reverse_dns)

    grouped = df.groupby(["date", "hostname"])["count"].sum().unstack(fill_value=0)

    ax = grouped.plot(kind="bar", stacked=True, figsize=(12, 6), colormap="tab20")
    plt.title("Email Volume by Source Host Group")
    plt.xlabel("Date")
    plt.ylabel("Email Count")

    plt.legend(
        title="Source",
        bbox_to_anchor=(1.05, 1),
        loc="upper left",
        borderaxespad=0.,
        frameon=True
    )

    plt.tight_layout()
    plt.show()

def main():
    import argparse
    parser = argparse.ArgumentParser(description="Visualize DMARC XML reports.")
    parser.add_argument("folder", help="Folder containing DMARC XML reports")
    args = parser.parse_args()

    df = load_all_reports(args.folder)
    if df.empty:
        print("No data found.")
        return

    plot_auth_results(df)
    plot_source_ip_volume(df)

if __name__ == "__main__":
    main()