#!/usr/bin/env python3
import socket
import signal
import requests

TRACE_ENABLE = "/sys/kernel/debug/tracing/events/mce/mce_record/enable"
TRACE_PIPE = "/sys/kernel/debug/tracing/trace_pipe"


def alert(alertname, host):
    dc = host[:3]
    alerts = [{
        "labels": {
            "alertname": alertname,
            "severity": "critical",
            "host": host,
            "team": "ceunix",
            "product": "all",
            "datacenter": dc,
        },
        "annotations": {
            "description": f"Machine Check Exception: Perhaps Correctable ECC. https://{host}-sp.{dc}.efscloud.net/",
        },
    }]
    requests.post(f"https://{dc}-prd-alerts0.{dc}.efscloud.net/api/v2/alerts", json=alerts)


def enable_trace():
    with open(TRACE_ENABLE, "w") as f:
        f.write("1")


def disable_trace():
    try:
        with open(TRACE_ENABLE, "w") as f:
            f.write("0")
    except Exception as e:
        print(f"[!] disable_trace() failed: {e}")
        pass


def main():
    host = socket.gethostname().split('.')[0]
    enable_trace()
    print("[*] MCE tracepoint enabled, reading trace_pipe...")

    def handle_exit(signum, frame):
        print("\n[!] Signal received, disabling tracepoint.")
        disable_trace()
        exit(0)

    signal.signal(signal.SIGINT, handle_exit)
    signal.signal(signal.SIGTERM, handle_exit)

    with open(TRACE_PIPE, "r") as trace:
        for line in trace:

            if "mce_record" not in line:
                continue

            alert("vmc.mce", host)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"[!] Exception: {e}")
    finally:
        disable_trace()
        print("[*] Tracepoint disabled.")
