# Copyright 2026 SigmaSRC, Inc. All Rights Reserved.

"""A persistent cache for Fuel code downloaded from the server.

The caching scheme works as follows:

- Module X.Y.Z version V is stored as $AVDATA/fuelcache/X/Y/Z.V.fuel.

- The most recently requested version of module X.Y.Z is stored as
  $AVDATA/fuelcache/X/Y/Z.fuel; this is a copy of
  $AVDATA/fuelcache/X/Y/Z.V.fuel.

- The first line of the cached file is a comment of the form
  # <fuel name="X.Y.Z" version="V" />


"""

import os
import re
import shutil
import http.client as httplib
import logging

from sigma.util import dance
from sigma.protocol import rest
from sigma.runtime.rtlib import RUNTIME_VERSION

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

INIT_CONTENTS = "# This file is necessary to make this directory a package.\n"

class FuelCache(object):

    def __init__(self, options):
        """Constructor.  Requires an Options argument."""
        self.options = options
        self.cachedir = options.datapath("fuelcache")
        self.table = {} # {module: [version, ...], ...}
        self.conn = None

    def add_module_versions(self, modules):
        """Adds a list of module versions.

        The argument should be the 'fuel' attribute of a Directive
        object; this is a list of Fuel elements each with a module and
        version attribute.

        This is used to later be able to retrieve the correct Fuel
        module.

        Note: it is possible to add multiple versions for the same
        module name; however in this case all subsequent calls to
        get_module() will fail unless an explicit version argument is
        given.

        """
        for mod in modules:
            versions = self.table.setdefault(mod.module, [])
            if mod.version not in versions:
                versions.append(mod.version)

    def get_module(self, name, version=None):
        """Gets a module from the cache.

        Returns the full pathname to the module if successful.

        Returns None if the module can't been downloaded.

        """
        if version is None:
            version = self._get_version(name)
            if version is None:
                return None
        basepath = name.replace(".", os.sep)
        relpath = basepath + ".fuel"
        fullpath = os.path.join(self.cachedir, relpath)
        if self._check_version(name, version, fullpath):
            return fullpath
        cachedrelpath = "%s.%d.fuel" % (basepath, version)
        cachedfullpath = os.path.join(self.cachedir, cachedrelpath)
        if os.path.exists(cachedfullpath):
            self._copy_file(cachedfullpath, fullpath)
            return fullpath
        text = self._get_module(name, version)
        if text is None:
            return None
        try:
            self._make_package(name)
        except (os.error, IOError) as err:
            logger.error("Can't create package %s: %s", name, err)
            return None
        newpath = cachedfullpath + ".new"
        try:
            f = open(newpath, "w", encoding="utf-8")
            try:
                f.write('# <fuel name="%s" version="%s" />\n' %
                        (name, version))
                f.write(text)
            finally:
                f.close()
        except IOError as err:
            logger.error("Can't write %s: %s", fullpath, err)
            return None
        dance.rename(newpath, cachedfullpath)
        self._copy_file(cachedfullpath, fullpath)
        return fullpath

    def _get_version(self, name):
        """Helper to get the version number declared for a module.

        Returns None if no version has been declared or more than one
        version has been declared.

        """
        versions = self.table.get(name)
        if not versions:
            logger.error("No known versions for module %s", name)
            return None
        if len(versions) > 1:
            logger.error("Ambiguous versions for module %s: %s",
                         name, versions)
            return None
        return versions[0]

    def _check_version(self, name, version, fullpath):
        """Helper to check if the requested version is already present.

        Returns True if the given full pathname is the correct version,
        False otherwise.

        """
        if not os.path.exists(fullpath):
            return False
        try:
            f = open(fullpath)
        except IOError as err:
            logger.warn("Module file %s exists but can't be read: %s",
                        fullpath, err)
            return False
        try:
            line = f.readline()
        finally:
            f.close()
        m = re.match(r'# <fuel name="(\w+)" version="(\d+)" />$', line)
        if m is None or m.group(1) != name:
            return False
        return int(m.group(2)) == version

    def _get_module(self, name, version):
        """Helper to retrieve the requested version of a module.

        Returns the module text if successful, None otherwise.

        """
        fresh = False
        if self.conn is None:
            self.conn = rest.Connection(ssl=False, options=self.options)
            fresh = True
        try:
            try:
                text = self.conn.getModule(RUNTIME_VERSION, name, version)
            except httplib.BadStatusLine:
                if fresh:
                    raise
                # Most likely the server timed out on us; retry once
                logger.warn("Can't get module %s version %s; retrying",
                            name, version)
                self.conn = rest.Connection(ssl=False, options=self.options)
                text = self.conn.getModule(RUNTIME_VERSION, name, version)
        except rest.ALLERRORS as err:
            self.conn = None
            logger.error("Can't get module %s version %s: %s: %s",
                         name, version, err.__class__.__name__, err)
            return None
        if text is None:
            logger.error("Didn't get module %s version %s", name, version)
            return None
        logger.info("Downloaded module %s version %s from server",
                    name, version)
        return text

    def _make_package(self, name):
        """Helper to create the package for the given module.

        This ensures both that the directories exist and that they
        contain a file named __init__.py.

        This may raise os.error and IOError.

        """
        dir = self.cachedir
        if not os.path.exists(dir):
            os.makedirs(dir)
        parts = name.split(".")
        del parts[-1]
        for part in parts:
            dir = os.path.join(dir, part)
            if not os.path.exists(dir):
                os.mkdir(dir)
            initfile = os.path.join(dir, "__init__.py")
            if not os.path.exists(initfile):
                f = open(initfile, "w")
                try:
                    f.write(INIT_CONTENTS)
                finally:
                    f.close()

    def _copy_file(self, srcpath, dstpath):
        """Helper to copy a file.

        This first copies to a temporary pathname with a .new extension,
        then atomically (if possible) renames that to the destination.

        """
        newdstpath = dstpath + ".new"
        shutil.copy(srcpath, newdstpath)
        err = dance.rename(newdstpath, dstpath)
        if err is not None:
            raise err
