# Copyright 2026 SigmaSRC, Inc. All Rights Reserved.

"""XML classes for the agent-server protocol."""

import os
import re
import logging
import datetime

from sigma.util import config, avalog, xmlparse
from sigma.util.xmlparse import ElementClass
from sigma.util.xmlparse import String, Boolean, Integer, Timestamp
from sigma.policies.xmldefs import Implementation
from sigma.sched.tasks import Task

options = config.Options()
logger = logging.getLogger("protocol.xmlproto")

class Cert(ElementClass):
    """The <cert> element is a sub-element of <credentials>."""
    __element__ = "cert"
    __characters__ = "text"

class PrivKey(ElementClass):
    """The <privKey> element is a sub-element of <credentials>."""
    __element__ = "privKey"
    __characters__ = "text"

class Credentials(ElementClass):
    """The <credentials> element is used by registration and cert renewal."""
    __element__ = "credentials"
    __attributes__ = {"agentID": Integer}
    __children__ = {Cert: "cert",
                    PrivKey: "privKey"}

class Remove(ElementClass):
    """Matrix <remove> element, a sub-element of <group>."""
    __element__ = "remove"
    __characters__ = "text"

class Add(ElementClass):
    """Matrix <add> element, a sub-element of <group>."""
    __element__ = "add"
    __characters__ = "text"

class Group(ElementClass):
    """Matrix <group> element, a sub-element of <matrix>."""
    __element__ = "group"
    __attributes__ = {"name": String}
    __children__ = {Add: "adds[]",
                    Remove: "removes[]"}

class Matrix(ElementClass):
    """The <matrix> element is used to represent the matrix (full and delta)."""
    __element__ = "matrix"
    __attributes__ = {"id": Integer,
                      "base": Integer}
    __children__ = {Group: "groups[]"}

class Upgrade(ElementClass):
    """The <upgrade> element is used by agent upgrade."""
    __element__ = "upgrade"
    __attributes__ = {"version": String,
                      "action": String,
                      "url": String,
                      "md5": String}

class Command(ElementClass):
    """The <command> element is written to $AVDATA/upgrades.xml."""
    __element__ = "command"
    __children__ = {
        Upgrade: "upgradeList[]",
        }

class Upload(ElementClass):
    """The <upload> element is used as an argument to uploadFile()."""
    __element__ = "upload"
    __attributes__ = {"fileID": Integer,
                      "type": String,
                      "startTime": Timestamp,
                      "stopTime": Timestamp,
                      "currentVersion": String,
                      "previousVersion": String,
                      "size": Integer}
    __characters__ = "text"

    def __init__(self, *args, **kwds):
        ElementClass.__init__(self, *args, **kwds)
        # Log cases where startTime > stopTime (might interest someone)
        if self.startTime is not None and self.stopTime is not None:
            assert isinstance(self.startTime, datetime.datetime)
            assert isinstance(self.stopTime, datetime.datetime)
            if self.startTime > self.stopTime:
                logger.warn("type=%s: startTime=%s > stopTime=%s",
                            self.type, self.startTime, self.stopTime)

# ---

class Value(ElementClass):
    "A rule actual parameter value."
    __element__ = "value"
    __attributes__ = {}
    __characters__ = "text"

class Parameter(ElementClass):
    "A rule actual parameter."
    __element__ = "parameter"
    __attributes__ = {"name": String,
                      "type": String,
                      "value": String}
    __children__ = {Value: "valueElement"}

class Rule(ElementClass):
    "A rule as found in a <directive> element."
    __element__ = "rule"
    __attributes__ = {"id": Integer,
                      "name": String,
                      "enforce": Boolean,
                      "category": String,
                      "enterprise": Boolean,
                      "template-id": Integer,
                      "template-version": Integer,
                      "module": String, # bw compat
                      "procedure": String, # bw compat
                     }
    __children__ = {Implementation: "implementations[]",
                    Parameter: "parameters[]",
                   }

class ShortDirective(ElementClass):
    "Represents the <directive> syntax found in the agent home page."
    __element__ = "directive"
    __attributes__ = {"id": Integer,
                      "version": Integer,
                     }

class Fuel(ElementClass):
    "Represents a Fuel module version declaration."
    __element__ = "fuel"
    __attributes__ = {"module": String,
                      "version": Integer,
                     }

class Directive(ShortDirective):
    "Represents the <directive> syntax found in full directives."
    __attributes__ = {"id": Integer,
                      "version": Integer,
                      "name": String,
                      "timestamp": Timestamp,
                      "enforce": Boolean,
                      "poliwall-mode": String, # "allow" or "deny"
                     }
    __children__ = {Rule: "rules[]",
                    Task: "task",
                    Fuel: "fuel[]",
                   }

class Exempt(ElementClass):
    "Used in the agent home page."
    __element__ = "exempt"
    __characters__ = "text"

    def get_rule_ids(self):
        return [int(x) for x in self.text.split(",") if x.strip()]

class Agent(ElementClass):
    "The agent home page, /Agent/agent.<agentID>.xml."
    __element__ = "agent"
    __attributes__ = {"id": Integer,
                      "enterprise": Boolean,
                      "poliwall-enforce": Boolean,
                      "certificate-fingerprint": String,
                     }
    __children__ = {ShortDirective: "directives[]",
                    Upgrade: "upgrades[]",
                    Exempt: "exempt",
                   }

def load_directive(key, options=None):
    filename = dir_filename(key, options)
    return Directive.__parseFile__(filename, strict=False)

policydir = options.datapath("policy")

def get_policydir(options=None):
    if options is not None:
        return options.datapath("policy")
    else:
        return policydir

def dir_filename(key, options=None):
    if os.sep in key or os.altsep and os.altsep in key or key.endswith(".xml"):
        return key
    pdir = get_policydir(options)
    return os.path.join(pdir, "dir.%s.xml" % key)

def create_policydir(options=None):
    pdir = get_policydir(options)
    if os.path.exists(pdir):
        return True
    try:
        os.makedirs(pdir)
        return True
    except os.error as err:
        logger.error("Can't create policy dir %s: %s", pdir, err)
        return False

def list_policydir(options=None):
    pdir = get_policydir(options)
    if os.path.exists(pdir):
        return os.listdir(pdir)
    else:
        return []

def get_policy_key(filename, options=None):
    m = re.match(r"\Adir\.(\w+\.\d+)\.xml\Z", filename)
    if m is not None:
        return m.group(1)
    else:
        return None
