# Copyright 2026 SigmaSRC, Inc. All Rights Reserved.

"""Utilities for working with the Sigma agent's credentials.

Implements a set of routines for accessing the agent's certificate,
id, and private key.

XXX To do:

- improve APIs to better support (re)registration

"""

# Python imports
import os
import ssl as _ssl
import hashlib
import time
import shutil
import http.client as httplib
import logging
import datetime

# Sigma imports
from sigma.util import avalog, config, dance


class _X509Cert(object):
    """Lightweight X509 certificate wrapper using stdlib ssl."""

    def __init__(self, pem_data):
        self._pem = pem_data
        self._der = _ssl.PEM_cert_to_DER_cert(pem_data)
        import tempfile
        f = tempfile.NamedTemporaryFile(mode='w', suffix='.pem', delete=False)
        f.write(self._pem)
        f.close()
        try:
            self._info = _ssl._ssl._test_decode_cert(f.name)
        finally:
            os.unlink(f.name)

    @classmethod
    def load_cert(cls, filename):
        with open(filename, 'r') as f:
            return cls(f.read())

    def get_serial_number(self):
        return int(self._info.get('serialNumber', '0'), 16)

    def get_not_after(self):
        return self._info.get('notAfter', '')

    def get_not_before(self):
        return self._info.get('notBefore', '')

    def get_subject(self):
        return self._info.get('subject', '')

    def get_issuer(self):
        return self._info.get('issuer', '')

    def get_version(self):
        return self._info.get('version', 0)

    def get_fingerprint(self, algorithm='sha512'):
        h = hashlib.new(algorithm)
        h.update(self._der)
        return h.hexdigest()

    def digest(self, algorithm='sha512'):
        return self.get_fingerprint(algorithm)

    def as_der(self):
        return self._der

# Global variables
logger = logging.getLogger("util.cert")
logger.info("Testing util.cert")

def safe_write(filename, data):
    """Safe write of data to disk.

    Used to write agent credentials to disk in a safe manner by
    creating a file with mode set to 0600.  Caller must be prepared to
    handle os.error exceptions from this routine.

    """

    fd = os.open(filename, os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o600)
    if isinstance(data, str):
        data = data.encode('utf-8')
    os.write(fd, data)
    os.fsync(fd)
    os.close(fd)

def file_load(filename):
    "Loads a credential file from disk returning None if it does not exist."
    try:
        if not os.path.isfile(filename):
            return None
        f = open(filename, "r")
        # these files will all be small
        val = f.read()
        f.close()
        return val
    except:
        logger.error("Exception attempting to read %s" % filename)
        return None

def fingerprint(x509_cert):
    """Returns the MD5 fingerprint of an x509 certificate

    Returns the MD5 fingerprint of an x509 certificate.  Outputs the same
    fingerprint as openssl x509 -fingerprint command.
    
    DEPRECATED in favor of running the cert.digest to output sha512.
    """
    # der = x509_cert.as_der()
    # md = md5.md5()
    # md.update(der)
    # hexdigest = md.hexdigest()
    hexdigest = x509_cert.get_fingerprint('sha512')
    hexlist = [hexdigest[x:x+2] for x in range(len(hexdigest)) if x%2 == 0]
    fingerprint = ':'.join(hexlist)
    return fingerprint

def print_cert(id, x509):
    """Print certificate.
    
    Prints a human readable summary of a certificate tagging it with id as
    a human readable identifier.

    """

    print("Certificate ", id)
    print("  Data")
    print("    Version: ", x509.get_version())
    print("    Serial Number: ", x509.get_serial_number())
    print("    Issuer: ", x509.get_issuer())
    print("    Validity")
    print("      Not before: ", x509.get_not_before())
    print("      Not After:  ", x509.get_not_after())
    print("    Subject: ", x509.get_subject())
    print("    Fingerprint: ", x509.digest('sha512'))

def fetch_cacert(http_server):
    "Fetches the CA certificate via http from the supplied url"
    conn = httplib.HTTPConnection(http_server)
    conn.request("GET", "/Certificate")
    resp = conn.getresponse()
    return resp.read(resp.length)

def get_expiration_date(x509):
    """Returns the expiration date of a certificate as a time.time() style
    float.

    """

    # get_not_after is some ASN1 native format, string extraction
    # seems to be the only way to get it out
    expire_date = str(x509.get_not_after())
    logger.info("Agent certificate expires on %s", expire_date)
    
    # convert the string from M2Crypto into a real date so we can check
    expire_date = time.strptime(expire_date,  "%b %d %H:%M:%S %Y %Z")
    expire_date = time.mktime(expire_date)

    return expire_date

class CertificateStore(object):
    "Provides access to the Sigma agent's credentials."

    def __init__(self, ssl=True, certdir=None, options=None):
        if certdir is None:
            if options is None:
                options = config.Options()
            certdir = options.agentcert
        self.certdir = certdir
        self.archivedir = os.path.join(self.certdir, "archive")
        self.idfile = os.path.join(self.certdir, "id.txt")
        if not os.path.exists(self.idfile):
            oldidfile = os.path.join(self.certdir, "id.avl")
            if os.path.exists(oldidfile):
                newidfile = self.idfile + ".new"
                try:
                    shutil.copyfile(oldidfile, newidfile)
                except (IOError, os.error) as err:
                    logger.error("Can't copy %s to %s: %s",
                                 oldidfile, newidfile, err)
                    self.idfile = oldidfile
                else:
                    if dance.rename(newidfile, self.idfile, logger) is None:
                        logger.info("Copied %s to %s", oldidfile, self.idfile)
                    else:
                        self.idfile = oldidfile
        self.certfile = os.path.join(self.certdir, "cert.pem")
        self.keyfile = os.path.join(self.certdir, "key.pem")
        self.cafile = os.path.join(self.certdir, "cacert.pem")

        agentID = file_load(self.idfile)
        if agentID is not None:
            agentID = int(agentID)
        self.agentID = agentID
        if ssl and os.path.isfile(self.certfile):
            self.cert = _X509Cert.load_cert(self.certfile)
        else:
            self.cert = None

    sslContext = None

    def getSSLContext(self):
        """Returns an SSL context suitably initialized for this cert store.

        Due to memory leaks in OpenSSL, we should reuse SSL contexts
        as much as possible.  Since the SSL context contains
        information derived from to the cert store, it is convenient
        to store the SSL context on the cert store.

        """
        if self.sslContext is not None:
            return self.sslContext
        import ssl as _ssl
        from sigma.protocol.ssl_callbacks import create_ssl_context
        ctx = create_ssl_context(cafile=self.cafile)
        self.sslContext = ctx
        self.fixSSLContext()
        return ctx

    def fixSSLContext(self):
        "Internal helper to update the SSL context once we have a cert."
        if (self.sslContext is not None and
            self.check_registration() and
            os.path.exists(self.certfile) and
            os.path.exists(self.keyfile)):
            self.sslContext.load_cert_chain(self.certfile, self.keyfile)

    def check_registration(self):
        """Checks whether the agent is already registered.

        This checks the on-disk idfile as well as the cached agent ID.

        """
        return self.agentID is not None and os.path.exists(self.idfile)

    def create_cert_dir(self):
        """Creates the certificate directory.

        Creates the certificate directory if it does not already exist.

        returns True on success; False when an os.error occurs (the
        error is also logged).

        """
        if not os.path.isdir(self.certdir):
            logger.info("Creating cert dir %r", self.certdir)
            try:
                os.mkdir(self.certdir, 0o700)
            except os.error as err:
                if not os.path.isdir(self.certdir):
                    logger.error("Failed to create cert dir %s: %s",
                                 self.certdir, err)
                    return False
        return True

    def create_archive_dir(self):
        """Creates the certificate archive directory.

        Creates a directory for archiving certificates if it does not
        already exist.

        returns True on success, False when an os.error occurs (error
        will be logged)

        """
        if not os.path.isdir(self.archivedir):
            logger.info("Creating cert archive dir: %r", self.archivedir)
            try:
                os.mkdir(self.archivedir, 0o700)
            except os.error as err:
                if not os.path.isdir(self.archivedir):
                    logger.error("Failed to create cert archive dir %s : %s",
                                 self.archivedir, err)
                    return False
        return True

    def cleanup_failed_registration(self):
        "Removes data files from disk if part of the registration failed."
        self.agentID = None
        self.cert = None
        oldidfile = os.path.join(self.certdir, "id.avl")

        try:
            if os.path.isfile(oldidfile):
                os.remove(oldidfile)
            if os.path.isfile(self.idfile):
                os.remove(self.idfile)
            if os.path.isfile(self.certfile):
                os.remove(self.certfile)
            if os.path.isfile(self.keyfile):
                os.remove(self.keyfile)

            logger.info("Cleaned up after failed registration")
        except:
            logger.error("Failed to clean up after failed registration")

    def write_id(self, agentID):
        "Writes a new agentID to disk.  Caller must handle exceptions."
        safe_write(self.idfile, "%s\n" % agentID)
        self.agentID = int(agentID)

    def write_cert(self, certificate):
        "Writes a new cert to disk.  Caller must handle exceptions."
        safe_write(self.certfile, certificate)
        self.cert = _X509Cert(certificate)

    def write_key(self, key):
        "Writes a new private key to disk.  Caller must handle exceptions."
        safe_write(self.keyfile, key)
        self.fixSSLContext()

    def write_cacert(self, cacert):
        "Writes a new CA certificate to disk.  Caller must handle exceptions."
        safe_write(self.cafile, cacert)

    def archive(self):
        """Archives the agent certificate

        Writes an archive copy of the current agent certificate to an archive
        directory without changing the current cert directory.  This method
        is used in prepairation for changing an agent's certificate.
        Certificates are archived using their md5 fingerprint plus the
        extension .cert or .key.

        """

        self.create_archive_dir()
        fprint = fingerprint(self.cert)
        filename = "".join(fprint.split(':'))
        archive_cert = os.path.join(self.archivedir, filename + ".cert")
        archive_key = os.path.join(self.archivedir, filename + ".key")
        if not os.path.isfile(archive_cert):
            try:
                # note always write cert and key together, just to be safe
                cert = file_load(self.certfile)
                key = file_load(self.keyfile)
                safe_write(archive_cert, cert)
                safe_write(archive_key, key)
                logger.info("Archived certificate %s", fprint)
            except os.error as err:
                logger.error("Unable to archive certificate: %s", err)
                return False
        return True

    def restore_from_archive(self, serialnum):
        """Restores a certificate from the archive."""

        if not os.path.isdir(self.archivedir):
            logger.error("No certificates are archived")
            return False

        oldcert = None
        for file in os.listdir(self.archivedir):
            if file.endswith('.cert'):
                file = os.path.join(self.archivedir, file)
                aCert = _X509Cert.load_cert(file)
                if aCert.get_serial_number() == serialnum:
                    oldcert = aCert
                    break
                
        # when we get to here oldcert is either None or the one we're
        # interested in
        if not oldcert:
            logger.error("Certificate %s not found", serialnum)
            return False

        fprint = fingerprint(oldcert).split(":")
        fprint = "".join(fprint)
        old_certfile = os.path.join(self.archivedir, fprint + ".cert")
        old_keyfile = os.path.join(self.archivedir, fprint + ".key")
        

        safe_write(self.certfile, file_load(old_certfile))
        safe_write(self.keyfile, file_load(old_keyfile))

        return True
