#! /usr/bin/env python3
# Copyright 2026 SigmaSRC, Inc. All Rights Reserved.

"""Policy Engine.

Usage: %prog [-p] <id>.<version>
   or: %prog [-p] <filename>

Execute all rules from the directive, and spool the compliance output
for upload to the server.  The argument is either of the form
<id>.<version>, which is converted to dir.<id>.<version>.xml in the
standard policy directory ($AVDATA/policy), or a filename from which
the directive is read (containing at least one slash or backslash or
ending in '.xml').

"""

# Python imports
import os
import sys
import copy
import glob
import pickle
import logging
import optparse

# Enter Python UTF-8 Mode before any I/O so non-ASCII host data / Fuel source is
# handled correctly regardless of host locale (re-execs once if not already on).
if __name__ == "__main__":
    from sigma.util.utf8 import ensure_utf8_mode
    ensure_utf8_mode()

# Sigma imports
import pwgather
from sigma.util import avalog, config, fslock, xmlparse, paramtypes
from sigma.util import compliance, ostypes
from sigma.util.enterprise import readEnterpriseFlag
from sigma.util.errors import usage
from sigma.spool import spooling
from sigma.fuelc import modreg, compile
from sigma.runtime import rtlib, pwall, xmlout, fuelcache
from sigma.protocol import xmlproto

# Global objects
options = config.Options()
logger = logging.getLogger("engine")
parser = optparse.OptionParser(usage=__doc__.strip())
parser.add_option("-p", "--poliwall", action="store_true", dest="poliwall",
                  help="special processing for poliwall rules")

def main(args=None):
    # Options parsing
    opts, args = parser.parse_args(args)
    if len(args) != 1:
        return usage("This program takes exactly one argument.")
    key = args[0]
    filename = xmlproto.dir_filename(key)

    # Set up logging
    avalog.setupLogging(options)

    # Read XML input
    try:
        directive = xmlproto.load_directive(key)
    except IOError as err:
        logger.error("Can't open directive %s: %s", key, err)
        return 1
    except xmlparse.ParseError as err:
        logger.error("ParseError: %s", err)
        return 1

    # Initialize the Fuel Cache
    cache = fuelcache.FuelCache(options)
    cache.add_module_versions(directive.fuel)

    # By default there are no poliwall rules
    pwrules = {}

    if not opts.poliwall:
        # Are there any poliwall rules in the directive?
        for rule in directive.rules:
            if rule.category in ("poliwall", "packet filter"):
                pwrules[rule.id] = [] # Assume compliance by default

        # If there are, run pwgather.py and get the pickle
        if pwrules:
            logger.info("Running pwgather.py")
            pwgather.main([])
            getpwcompliance(pwrules, directive.id)

    # Lock out other policy engines; give up after 30 minutes
    # XXX Should we continue without the lock if we fail to lock?
    # XXX What's a good timeout here?
    lockdir = options.datapath("policy.lock")
    logger.info("Locking %s", lockdir)
    lock = fslock.FSLock(lockdir)
    if not lock.lock(30*60):
        logger.error("Giving up on locking %s", lockdir)
        return 1
    logger.info("Acquired lock %s", lockdir)
    if not os.path.exists(filename):
        logger.error("Directive file %r disappeared", filename)
        logger.info("Releasing lock %s", lockdir)
        lock.unlock()
        return 1

    # Read exemptions
    exempt = read_exemptions()

    # Run the rules
    logger.info("Directive %s (enforce=%s)", filename, directive.enforce)
    host = rtlib.Host("localhost")
    host.pw_update_commands = opts.poliwall
    qitem = spooling.new(".xml")
    f = open(qitem.filename, "w", encoding="utf-8")
    host.xmlout = xmlout.ResultsOut(f)
    host.xmlout.results(directive.id, directive.version)
    if opts.poliwall:
        host.setruleid(None)
        host.setdirectiveid(None, None)
        pw = host.poliwall()
        pw.clear()
        pw.save()
    sts = run_rules(directive, host, pwrules, exempt, cache)
    if opts.poliwall:
        host.setruleid(None)
        host.setdirectiveid(None, None)
        pw = host.poliwall()
        pw.save_directives()
    host.xmlout.close()
    size = f.tell()
    f.close()
    if opts.poliwall:
        qitem.discard()
        logger.debug("Not spooling dummy poliwall compliance output (-p)")
    elif size and directive.id and directive.version:
        deployment_id = "%s.%s" % (directive.id, directive.version)
        qitem.spool(type="compliance", deployment_id=deployment_id)
        logger.info("Spooled %s", qitem.filename)
    else:
        qitem.discard()
        logger.warning("No compliance output, not spooling")
    logger.info("Releasing lock %s", lockdir)
    lock.unlock()
    return 0

def read_exemptions():
    fn = options.datapath("exempt.txt")
    try:
        f = open(fn)
    except IOError as err:
        logger.warn("Can't read %s: %s", fn, err)
        return set()
    ruleids = set()
    try:
        lineno = 0
        for line in f:
            lineno += 1
            try:
                ruleid = int(line)
            except ValueError as err:
                logger.warn("Bad value in file %s, line %d: %s",
                            fn, lineno, ruleid)
            else:
                ruleids.add(ruleid)
    finally:
        f.close()
    return ruleids

def getpwenforce():
    if not options.poliwall_enforce:
        logger.info("Poliwall enforcement: disabled by configuration option")
        return False
    pwstatefile = options.datapath("pwstate.txt")
    try:
        f = open(pwstatefile)
    except IOError as err:
        logger.info("Poliwall enforcement: disabled by IOError: %s", err)
        return False
    kernel = False
    try:
        for line in f:
            if line.startswith("kernel="):
                kernel = line[7:].strip().lower() == "true"
    finally:
        f.close()
    if kernel:
        logger.info("Poliwall enforcement: ENABLED")
    else:
        logger.info("Poliwall enforcement: disabled")
    return kernel

def getpwcompliance(pwrules, directive_id):
    lock = pwgather.getlock()
    if lock is None:
        return
    try:
        pwcompliance = pwgather.readpwcompliance()
        logger.debug("Before: pwrules = %r", pwrules)
        d = pwcompliance.get(directive_id, {})
        for key in list(d.keys()):  # list(): we delete from d while iterating (Py3)
            if key not in pwrules:
                # This rule is no longer in the directive
                del d[key]
            elif d[key]:
                pwrules[key] = d[key]
                d[key] = []
        for key in pwrules:
            if key not in d:
                d[key] = []
        pwcompliance[directive_id] = d
        logger.debug("After: pwrules = %r", pwrules)
        # XXX This may lose info if we crash before the spool() call
        pwgather.writepwcompliance(pwcompliance)
    finally:
        pwgather.unlock(lock)

def run_rules(directive, host, pwrules, exempt, cache):
    total_rules = 0
    total_compliant = 0
    total_noncompliant = 0

    if directive.poliwall_mode is not None:
        if directive.poliwall_mode.lower() in pwall.AllowedPoliwallModes:
            host.pwall_mode = directive.poliwall_mode.lower()
            logger.info("Poliwall mode: %s", host.pwall_mode)
        else:
            logger.error("Bogus poliwall-mode attribute (%r) on directive",
                         directive.poliwall_mode)
            host.pwall_mode = None

    pwall_enforce = getpwenforce()

    cm = compile.CompilationManager(options.fuelpath)
    save_sys_path = list(sys.path)
    sts = 0
    host.context = copy.copy(directive)
    host.context.rules = []
    ostype = ostypes.guess_ostype()
    enterprise = readEnterpriseFlag()

    # Main loop
    for rule in directive.rules:
        total_rules += 1

        # Set the rule context for actions.logaction()
        host.context.rules = [rule]

        # Calculate effective enforce flag from directive.enforce
        # and rule.enforce.
        enforce = (directive.enforce and rule.enforce != False or
                   rule.enforce and directive.enforce != False)

        # Then set host.enforce and host.pwall_enforce accordingly;
        # these also take configuration options and the poliwall
        # kernel mode into account.
        if not enforce:
            host.enforce = host.pwall_enforce = False
        else:
            host.enforce = options.host_enforce
            host.pwall_enforce = pwall_enforce

        # Set rule id early to aid diagnostics
        host.setruleid(rule.id)
        host.setdirectiveid(directive.id, directive.version)

        # Filter out exempt rules
        if rule.id in exempt:
            logger.info("Rule %s: exempt from execution", rule.id)
            host.xmlout.leave("Skipped: rule is exempt on this agent")
            sts = 1
            continue

        # Filter out enterprise rules
        if rule.enterprise and not enterprise:
            logger.info("Rule %s: enterprise only", rule.id)
            host.xmlout.leave("Skipped: not an enterprise agent")
            sts = 1
            continue

        # Extract the module and procedure name
        module, procedure = rule.module, rule.procedure
        version = None
        if module is None or procedure is None:
            module, procedure, version = find_best_implementation(rule, ostype)
            if module is None or procedure is None:
                logger.info("Rule %s: unimplemented on %s", rule.id, ostype)
                host.xmlout.leave("Skipped: not applicable on %s" % ostype)
                sts = 1
                continue

        # Log the rule id, Fuel name, and effective enforcement flags
        logger.info("Rule %s: %s.%s (enforce=%s, pwall_enforce=%s)",
                    rule.id, module, procedure,
                    host.enforce, host.pwall_enforce)

        # Set the cache depending on whether the module is cached or not
        if module.startswith("pw.") or rule.category == "packet filter":
            # Never get poliwall code from the cache
            fn = None
        else:
            # See if the cache has it (or can get it)
            fn = cache.get_module(module)
        if fn:
            logger.info("Using cached Fuel for %s version %s", module, version)
            fuelpath = [cache.cachedir] + options.fuelpath
        else:
            logger.info("Using preloaded Fuel for %s", module)
            fuelpath = options.fuelpath
        cm.setfuelpath(fuelpath)
        sys.path = save_sys_path + fuelpath

        # Load the module
        try:
            defn = cm.importmodule(module)
        except modreg.LoadError as err:
            logger.error("Rule %s: %s", rule.id, err)
            host.xmlout.error("Fuel import error: %s" % err)
            sts = 1
            continue

        # Compute arguments and check entry point
        args = [host]
        argtuples = []
        # The <parameter> elements must appear in order
        errors = False
        for param in rule.parameters:
            value = None
            if param.valueElement is not None:
                value = param.valueElement.text
            if value is None:
                value = param.value
            if value:
                try:
                    value = paramtypes.convert(value, param.type)
                except paramtypes.ParamError as err:
                    logger.error("Failed to convert value %r "
                                 "to type %r for parameter %r (%s: %s)" %
                                 (value, param.type, param.name,
                                  err.__class__.__name__, err))
                    value = None
                    errors = True
            else:
                value = None
            args.append(value)
            # Do this after the value has been converted -- shouldn't fail
            argtuples.append((param.name, paramtypes.getalias(param.type)))
        if errors:
            host.xmlout.error("Fuel check error: %s" % err)
            sts = 1
            continue
        logger.debug("args = %r", args)
        try:
            cm.check_entrypoint(defn.value, procedure, argtuples)
        except modreg.LoadError as err:
            logger.error("Rule %s (%s.%s): %s",
                         rule.id, module, procedure, err)
            host.xmlout.error("Fuel check error: %s" % err)
            sts = 1
            continue

        # Import the Python code
        modname = module + "_fuel"
        parts = modname.split(".")
        try:
            mod = __import__(modname)
            for attr in parts[1:]:
                mod = getattr(mod, attr)
            proc = getattr(mod, procedure)
        except KeyboardInterrupt:
            raise
        except Exception as err:
            logger.error("Rule %s: can't import %s from %s: %s",
                         rule.id, procedure, modname, err)
            host.xmlout.error("Python import error: %s: %s" %
                              (err.__class__.__name__, err))
            sts = 1
            continue

        # Execute the rule
        host.settmplid(rule.template_id)
        msgs = pwrules.get(rule.id)
        if msgs:
            for msg in msgs:
                # Note: format_compliance() in pwgather.py puts
                # "blocked" in the message for blocked packets and
                # "seen" for packets that were let through
                if "blocked" in msg:
                    host.xmlout.enforcement(msg)
                else:
                    host.xmlout.violation(msg)
        try:
            proc(*args)
        except rtlib.FuelExit as err:
            logger.info("Rule %s: FuelExit(%d)", rule.id, err.status)
            if err.status:
                sts = 1
        except KeyboardInterrupt:
            raise
        except Exception as err:
            logger.exception("Rule %s: run-time exception", rule.id)
            host.xmlout.error("Run-time exception (%s.%s: %s)" %
                              (err.__class__.__module__,
                               err.__class__.__name__,
                               str(err)))
            sts = 1
        if host.xmlout.violations:
            logger.warning("Rule %s: compliance violation", rule.id)
            total_noncompliant += 1
        else:
            if host.xmlout.status() == 0:
                total_compliant += 1
        if host.xmlout.enforcements:
            logger.warning("Rule %s: enforcement action taken", rule.id)

    if not host.pw_update_commands:
        key = "%s.%s" % (directive.id, directive.version)
        total_n_a = total_rules - total_compliant - total_noncompliant
        c = compliance.Compliance(options)
        c.set(key, total_rules, total_compliant, total_n_a, directive.name)
        c.save()

    # After the for loop is over: return 0 for success, 1 for errors
    return sts

def find_best_implementation(rule, ostype):
    if hasattr(rule, "implementations"):
        for ost in ostypes.allancestors(ostype):
            for impl in rule.implementations:
                ostypes_list = impl.ostypes.split(",")
                if ost in ostypes_list:
                    return impl.module, impl.procedure, impl.version
    return None, None, None

if __name__ == "__main__":
    try:
        sys.exit(main())
    except (SystemExit, KeyboardInterrupt):
        raise
    except:
        avalog.setupLogging()
        logger.exception("Unexpected exception:")
