# Copyright 2026 SigmaSRC, Inc. All Rights Reserved.

"""Matrix support.

There are two matrix IDs with special meanings: id 0 is interpreted as
``not interested in the matrix''; id 1 is the initial, completely
empty, matrix.

"""

import os
import logging
import pickle

from sigma.util import config, dance, fuelset
from sigma.runtime.pwall import dotted2numeric, IPAddr, RangeSet
from sigma.protocol import xmlproto

logger = logging.getLogger("runtime.matrix")

def get_matrixfn(options=None):
    if options is None:
        options = config.Options()
    return options.datapath("matrix.pickle")

def parse_range(text):
    """Helper function to parse the body text of <add> and <remove>."""
    result = RangeSet()
    if not text or text.isspace():
        return result
    text = str(text) # We don't need no stinkin' Unicode here
    pairs = [x.strip() for x in text.split(",") if not x.isspace()]
    for pair in pairs:
        if "-" in pair:
            lo, hi = pair.split("-", 1)
            result.add(dotted2numeric(lo.strip()), dotted2numeric(hi.strip()))
        else:
            result.add(dotted2numeric(pair))
    return result

class Matrix(object):

    def __init__(self, options=None):
        self.matrixfn = get_matrixfn(options)
        self.id = 1
        self.groups = {}

    def gethostgroup(self, name):
        grp = self.groups.get(name)
        if grp is None:
            return fuelset.FuelSet()
        if isinstance(grp, RangeSet):
            return grp.map(IPAddr)
        return fuelset.FuelSet(map(IPAddr, grp))

    def gethostgroupnames(self):
        return fuelset.FuelSet(self.groups.keys())

    def getid(self):
        return self.id

    def load(self):
        if not os.path.exists(self.matrixfn):
            return False
        try:
            f = open(self.matrixfn, "rb")
        except IOError as msg:
            logger.error("Can't read %s: %s", self.matrixfn, msg)
            return False
        try:
            try:
                id = pickle.load(f)
                groups = pickle.load(f)
            except Exception as msg:
                logger.error("Can't load pickle from %s: %s",
                             self.matrixfn, msg)
                return False
            else:
                self.id, self.groups = id, groups
                logger.info("Loaded matrix id=%r", self.id)
                return True
        finally:
            f.close()

    def store(self):
        tfn = self.matrixfn + ".new"
        try:
            f = open(tfn, "wb")
        except IOError as msg:
            logger.error("Can't write %s: %s", tfn, msg)
            return
        try:
            pickle.dump(self.id, f, pickle.HIGHEST_PROTOCOL)
            pickle.dump(self.groups, f, pickle.HIGHEST_PROTOCOL)
        finally:
            f.close()
        if dance.rename(tfn, self.matrixfn, logger) is None:
            logger.info("Stored matrix id=%r", self.id)

    def update(self, update):
        assert isinstance(update, xmlproto.Matrix), repr(update)
        assert isinstance(update.id, int)
        assert isinstance(update.base, int)
        assert isinstance(self.id, int)
        if update.base == 1:
            logger.info("Got full matrix (id=%r)", update.id)
            self.groups = {}
        elif update.base == self.id:
            logger.info("Got matrix delta (base=%r, id=%r)",
                        update.base, update.id)
        else:
            logger.error("Got bogus matrix delta (base=%r, id=%r), resetting",
                        update.base, update.id)
            self.groups = {}
            self.id = 1
            return False
        for group in update.groups:
            assert isinstance(group.name, str), repr(group.name)
            mygroup = self.groups.get(group.name, RangeSet())
            for add in group.adds:
                mygroup |= parse_range(add.text)
            for remove in group.removes:
                mygroup -= parse_range(remove.text)
            self.groups[group.name] = mygroup
        self.id = update.id
        return True

def loadmatrix(options=None):
    m = Matrix(options)
    m.load()
    return m

def getmatrixid(options=None):
    id = 1
    matrixfn = get_matrixfn(options)
    try:
        f = open(matrixfn, "rb")
    except IOError as err:
        logger.debug("Can't open %s: %s", matrixfn, err)
    else:
        try:
            id = pickle.load(f)
            id = int(id)
        except Exception as err:
            logger.debug("Can't parse matrix id: %s", err)
        f.close()
    logger.debug("Matrix id=%r", id)
    return id
