import ssl
import socket
from datetime import datetime, timezone
import os
import logging
from kubernetes import client, config
import tempfile

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s"
)

DEPLOYMENT = os.getenv("DEPLOYMENT_NAME", "my-deployment")
NAMESPACE = os.getenv("NAMESPACE", "default")
SECRET_CERT_PATH = os.getenv("SECRET_CERT_PATH", "/certs/tls.crt")
SERVICE_HOST = os.getenv("SERVICE_HOST", "my-app.default.svc")
SERVICE_PORT = int(os.getenv("SERVICE_PORT", 443))

def read_secret_cert():
    logging.info(f"Reading certificate from secret path: {SECRET_CERT_PATH}")
    return ssl._ssl._test_decode_cert(SECRET_CERT_PATH)

def get_secret_cert_expiry():
    cert = read_secret_cert()
    expiry = datetime.strptime(cert['notAfter'], "%b %d %H:%M:%S %Y %Z")
    logging.info(f"Secret certificate expires at: {expiry}")
    return expiry

def get_common_name(cert):
    """
    Extract Common Name (CN) from a cert dict as returned by ssl.getpeercert()
    or ssl._ssl._test_decode_cert(). Returns None if not found.
    """
    if not cert:
        return None
    subject = cert.get('subject', [])
    for rdn in subject:
        # rdn is a tuple of tuples, e.g. (('commonName', 'example.com'),)
        for attr in rdn:
            if isinstance(attr, tuple) and len(attr) == 2:
                key, val = attr
                if key.lower() in ('commonname', 'cn'):
                    return val
    return None

def get_live_cert_expiry():
    logging.info(f"Connecting to service {SERVICE_HOST}:{SERVICE_PORT} to get live certificate")
    # Try the secure default context first (will verify hostname)
    ctx = ssl.create_default_context()
    try:
        with ctx.wrap_socket(socket.socket(), server_hostname=SERVICE_HOST) as s:
            s.connect((SERVICE_HOST, SERVICE_PORT))
            # request both the dict view and the binary DER form
            cert = s.getpeercert()  # may be empty
            der = s.getpeercert(binary_form=True)
    except ssl.SSLCertVerificationError as e:
        logging.warning("Certificate verification failed (%s). Falling back to non-verifying connection to read cert (insecure).", e)
        # Insecure fallback: disable hostname and certificate verification just to read the cert
        ctx = ssl.create_default_context()
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE
        with ctx.wrap_socket(socket.socket(), server_hostname=SERVICE_HOST) as s:
            s.connect((SERVICE_HOST, SERVICE_PORT))
            cert = s.getpeercert()
            der = s.getpeercert(binary_form=True)

    # If the dict view lacks 'notAfter', try decoding the DER to get a full dict
    if not cert or 'notAfter' not in cert:
        if der:
            pem = ssl.DER_cert_to_PEM_cert(der)
            # write PEM to a temp file and decode using ssl test decoder to get 'notAfter'
            with tempfile.NamedTemporaryFile(mode="w", delete=True) as tf:
                tf.write(pem)
                tf.flush()
                cert = ssl._ssl._test_decode_cert(tf.name)
        else:
            raise RuntimeError("Failed to obtain certificate from peer")

    expiry = datetime.strptime(cert['notAfter'], "%b %d %H:%M:%S %Y %Z")
    logging.info(f"Live certificate expires at: {expiry}")
    return expiry, cert

def restart_deployment():
    logging.info(f"Restarting deployment {DEPLOYMENT} in namespace {NAMESPACE}")
    config.load_incluster_config()
    apps = client.AppsV1Api()
    body = {
        "spec": {
            "template": {
                "metadata": {
                    "annotations": {
                        "cert-checker/restartedAt": datetime.now(timezone.utc).isoformat()
                    }
                }
            }
        }
    }
    apps.patch_namespaced_deployment(
        name=DEPLOYMENT,
        namespace=NAMESPACE,
        body=body
    )
    logging.info(f"Deployment {DEPLOYMENT} restarted")

if __name__ == "__main__":
    logging.info("Starting certificate check")
    secret_cert = read_secret_cert()
    secret_expiry = datetime.strptime(secret_cert['notAfter'], "%b %d %H:%M:%S %Y %Z")
    live_expiry, live_cert = get_live_cert_expiry()

    secret_cn = get_common_name(secret_cert)
    live_cn = get_common_name(live_cert)

    logging.info(f"Secret expiry: {secret_expiry}")
    logging.info(f"Live cert expiry: {live_expiry}")
    logging.info(f"Secret CN: {secret_cn}")
    logging.info(f"Live CN: {live_cn}")

    if secret_cn != live_cn:
        logging.warning("Certificate mismatch detected (CN), check your setup.")
        exit(0)

    if secret_expiry != live_expiry:
        logging.warning("Certificate mismatch detected (expiry), restarting deployment")
        restart_deployment()
    else:
        logging.info("Certificates match, nothing to do.")
