HPI/my/core/structure.py

from __future__ import annotations

import atexit
import os
import shutil
import sys
import tarfile
import tempfile
import zipfile
from collections.abc import Generator, Sequence
from contextlib import contextmanager
from pathlib import Path

from .logging import make_logger

logger = make_logger(__name__, level="info")


def _structure_exists(base_dir: Path, paths: Sequence[str], *, partial: bool = False) -> bool:
    """
    Helper function for match_structure to check if
    all subpaths exist at some base directory

    For example:

    dir1
    ├── index.json
    └── messages
        └── messages.csv

    _structure_exists(Path("dir1"), ["index.json", "messages/messages.csv"])
    """
    targets_exist = ((base_dir / f).exists() for f in paths)
    if partial:
        return any(targets_exist)
    else:
        return all(targets_exist)


ZIP_EXT = {".zip"}
TARGZ_EXT = {".tar.gz"}


@contextmanager
def match_structure(
    base: Path,
    expected: str | Sequence[str],
    *,
    partial: bool = False,
) -> Generator[tuple[Path, ...], None, None]:
    """
    Given a 'base' directory or archive (zip/tar.gz), recursively search for one or more paths that match the
    pattern described in 'expected'. That can be a single string, or a list
    of relative paths (as strings) you expect at the same directory.

    If 'partial' is True, it only requires that one of the items in
    expected be present, not all of them.

    This reduces the chances of the user misconfiguring gdpr exports, e.g.
    if they archived the folders instead of the parent directory or vice-versa

    When this finds a matching directory structure, it stops searching in that subdirectory
    and continues onto other possible subdirectories which could match

    If base is an archive, this extracts it into a temporary directory
    (configured by core_config.config.get_tmp_dir), and then searches the extracted
    folder for matching structures

    This returns the top of every matching folder structure it finds

    As an example:

    export_dir
    ├── exp_2020
    │   ├── channel_data
    │   │   ├── data1
    │   │   └── data2
    │   ├── index.json
    │   ├── messages
    │   │   └── messages.csv
    │   └── profile
    │       └── settings.json
    └── exp_2021
        ├── channel_data
        │   ├── data1
        │   └── data2
        ├── index.json
        ├── messages
        │   └── messages.csv
        └── profile
            └── settings.json

    Giving the top directory as the base, and some expected relative path like:

    with match_structure(Path("export_dir"), expected=("messages/messages.csv", "index.json")) as results:
        # results in this block is (Path("export_dir/exp_2020"), Path("export_dir/exp_2021"))

    This doesn't require an exhaustive list of expected values, but its a good idea to supply
    a complete picture of the expected structure to avoid false-positives

    This does not recursively decompress archives in the subdirectories,
    it only unpacks into a temporary directory if 'base' is an archive

    A common pattern for using this might be to use get_files to get a list
    of archives or top-level gdpr export directories, and use match_structure
    to search the resulting paths for an export structure you're expecting
    """
    from . import core_config as CC

    tdir = CC.config.get_tmp_dir()

    if isinstance(expected, str):
        expected = (expected,)

    is_zip: bool = base.suffix in ZIP_EXT
    is_targz: bool = any(base.name.endswith(suffix) for suffix in TARGZ_EXT)

    searchdir: Path = base.absolute()
    try:
        # if the file given by the user is an archive, create a temporary
        # directory and extract it to that temporary directory
        #
        # this temporary directory is removed in the finally block
        if is_zip or is_targz:
            # sanity check before we start creating directories/rm-tree'ing things
            assert base.exists(), f"archive at {base} doesn't exist"

            searchdir = Path(tempfile.mkdtemp(dir=tdir))

            if is_zip:
                # base might already be a ZipPath, and str(base) would end with /
                zf = zipfile.ZipFile(str(base).rstrip('/'))
                zf.extractall(path=str(searchdir))
            elif is_targz:
                with tarfile.open(str(base)) as tar:
                    # filter is a security feature, will be required param in later python version
                    mfilter = {'filter': 'data'} if sys.version_info[:2] >= (3, 12) else {}
                    tar.extractall(path=str(searchdir), **mfilter)  # type: ignore[arg-type]
            else:
                raise RuntimeError("can't happen")
        else:
            if not searchdir.is_dir():
                raise NotADirectoryError(f"Expected either a zip/tar.gz archive or a directory, received {searchdir}")

        matches: list[Path] = []
        possible_targets: list[Path] = [searchdir]

        while len(possible_targets) > 0:
            p = possible_targets.pop(0)

            if _structure_exists(p, expected, partial=partial):
                matches.append(p)
            else:
                # extend the list of possible targets with any subdirectories
                for f in os.scandir(p):
                    if f.is_dir():
                        possible_targets.append(p / f.name)

        if len(matches) == 0:
            logger.warning(f"""While searching {base}, could not find a matching folder structure. Expected {expected}. You're probably missing required files in the gdpr/export""")

        yield tuple(matches)

    finally:

        if is_zip or is_targz:
            # make sure we're not mistakenly deleting data
            assert str(searchdir).startswith(str(tdir)), f"Expected the temporary directory for extracting archive to start with the temporary directory prefix ({tdir}), found {searchdir}"

            shutil.rmtree(str(searchdir))


def warn_leftover_files() -> None:
    from . import core_config as CC

    base_tmp: Path = CC.config.get_tmp_dir()
    leftover: list[Path] = list(base_tmp.iterdir())
    if leftover:
        logger.debug(f"at exit warning: Found leftover files in temporary directory '{leftover}'. this may be because you have multiple hpi processes running -- if so this can be ignored")


atexit.register(warn_leftover_files)