gather.py

"""
Methods for gathering data.

Gathers information about all puppet modules, including which are
present in our environment, their metadata, and their output of
``puppet strings``.
"""

from dataclasses import dataclass
from typing import (
    Any,
)
import json
import os.path
import hashlib
from .puppet.strings import puppet_strings
from .cache import Cache


@dataclass
class ModuleEntry:
    """
    One entry in a module.

    Parameters:
        name - local name of the module, should always be the basename
               of path
        path - Absolute path in the filesystem where the module can be
               found.
        strings_output - output of `puppet strings`.
    """

    name: str
    path: str
    strings_output: bytes
    metadata: dict[str, Any]

    def file(self, path: str) -> str:
        """Return the absolute path of a path inside the module."""
        return os.path.join(self.path, path)


def get_puppet_strings(cache: Cache, path: str) -> bytes:
    """
    Run puppet string, but check cache first.

    The cache uses the contents of metadata.json as its key,
    so any updates without an updated metadata.json wont't be
    detected.

    Hashing the entire contents of the module was tested, but was to
    slow.
    """
    try:
        with open(os.path.join(path, 'metadata.json'), 'rb') as f:
            data = f.read()
            key = 'puppet-strings' + hashlib.sha1(data).hexdigest()
            if parsed := cache.get(key):
                result = parsed
            else:
                result = puppet_strings(path)
                cache.put(key, result)
            return result
    except FileNotFoundError:
        # TODO actually run puppet strings again.
        # This is just since without a metadata.json we always get a
        # cache miss, which is slow.
        # return puppet_strings(path)
        return b''

        # try:
        #     with open(module.file('.git/FETCH_HEAD')) as f:
        #         st = os.stat(f.fileno())
        #         st.st_mtime
        # except FileNotFoundError:
        #     pass


def get_modules(cache: Cache, dir: str) -> list[ModuleEntry]:
    """
    Enumerate modules in directory.

    The directory should be the modules subdirectory of an environment,
    e.g. /etc/puppetlabs/code/environments/production/modules.
    """
    modules: list[ModuleEntry] = []

    for entry in sorted(list(os.scandir(dir)), key=lambda d: d.name):
        # TODO Logging
        # print('- entry', entry, file=sys.stderr)
        name = entry.name
        path = os.path.join(dir, entry)
        strings_data = get_puppet_strings(cache, path)

        try:
            with open(os.path.join(path, 'metadata.json')) as f:
                metadata = json.load(f)
        except FileNotFoundError:
            metadata = {}

        modules.append(ModuleEntry(name, path, strings_data, metadata))

    return modules