<commit_before>"""Serve nix store objects over HTTP."""
import argparse
import logging
import os
from os.path import exists, isdir, join, basename, dirname
import re
from subprocess import check_output, Popen, PIPE, CalledProcessError
from threading import Thread

from flask import Flask, make_response, send_file, request, jsonify
import six

from servenix import __version__
from servenix.common.utils import decode_str, strip_output, find_nix_paths
from servenix.common.exceptions import (NoSuchObject, NoNarGenerated,
                                        BaseHTTPError, NixImportFailed,
                                        CouldNotUpdateHash, ClientError)

_HASH_REGEX=re.compile(r"[a-z0-9]{32}")
_PATH_REGEX=re.compile(r"([a-z0-9]{32})-[^' \n/]*$")

class NixServer(Flask):
    """Serves nix packages."""
    def __init__(self, nix_store_path, nix_bin_path, nix_state_path,
                 compression_type, debug):
        # Path to the local nix store.
        self._nix_store_path = nix_store_path
        # Path to the local nix state directory.
        self._nix_state_path = nix_state_path
        # Path to the folder containing nix binaries.
        self._nix_bin_path = nix_bin_path
        # Matches valid nix store paths (local to this store)
        self._full_store_path_regex = re.compile(
            join(self._nix_store_path, _PATH_REGEX.pattern))
        self._compression_type = compression_type
        # Cache mapping object hashes to store paths.
        self._hashes_to_paths = {}
        # Cache mapping object hashes to store paths which have been validated.
        self._hashes_to_valid_paths = {}
        # Cache mapping store paths to object info.
        self._paths_to_info = {}
        # Set of known store paths.
        self._known_store_paths = set()
        # A static string telling a nix client what this store serves.
        self._cache_info = "\n".join([
            "StoreDir: {}".format(self._nix_store_path),
            "WantMassQuery: 1",
            "Priority: 30"
        ]) + "\n"
        if self._compression_type == "bzip2":
            self._content_type = "application/x-bzip2"
            self._nar_extension = ".nar.bz2"
        else:
            self._content_type = "application/x-xz"
            self._nar_extension = ".nar.xz"
        # Enable interactive debugging on unknown errors.
        self._debug = debug

    def store_path_from_hash(self, store_object_hash):
        """Look up a store path using its hash.

        The name of every object in the nix store starts with a
        32-character hash. We can find the full path by finding an
        object that starts with this hash.

        :param store_object_hash: The 32-character hash prefix of the object.
        :type store_object_hash: ``str``

        :return: The full store path to the object.
        :rtype: ``str``

        :raises: :py:class:`NoSuchObject` if the object isn't in the store.
        """
        if store_object_hash in self._hashes_to_valid_paths:
            # Case: the hash maps to a path which has been validated
            # to exist in the nix store and in the database.
            return self._hashes_to_valid_paths[store_object_hash]
        elif store_object_hash in self._hashes_to_paths:
            # Case: the hash maps to a path which is known to exist in
            # the nix-store, but hasn't been checked for existence in
            # the database. Check that it's in the DB before returning.
            path = self._hashes_to_paths[store_object_hash]
            if self.check_in_store(path):
                self._hashes_to_valid_paths[store_object_hash] = path
                return path
            # Otherwise, remove it from the known hashes and raise an error.
            self._hashes_to_paths.pop(store_object_hash, None)
            raise NoSuchObject("No object with hash {} was found!"
                               .format(store_object_hash))
        # Get the list of store objects by listing the directory.
        # Iterate through them until a matching hash is found, or
        # we've exhausted all paths, in which case we error.
        for path in map(decode_str, os.listdir(self._nix_store_path)):
            match = _PATH_REGEX.match(path)
            if match is None:
                continue
            path = join(self._nix_store_path, path)
            prefix = match.group(1)
            # Add every path seen to the _hashes_to_paths cache.
            self._hashes_to_paths[prefix] = path
            if prefix == store_object_hash:
                # The path exists in the store. Ensure it's also a
                # valid path according to nix-store.
                if not self.check_in_store(path):
                    break
                self._hashes_to_valid_paths[store_object_hash] = path
                return path
        raise NoSuchObject("No object with hash {} was found."
                           .format(store_object_hash))

    def check_in_store(self, store_path):
        """Check that a store path exists in the nix store.

        :param store_path: Path to an object in the nix store. Is assumed
            to match the full path regex.
        :type store_path: ``str``

        :return: True if it's in the store, and False otherwise.
        :rtype: ``bool``
        """
        if store_path in self._known_store_paths:
            return True
        elif not exists(store_path):
            # If the path isn't in the filesystem, it definitely is
            # not a valid path.
            return False
        try:
            # If it is on the filesystem, it doesn't necessarily mean
            # that it's a registered path in the store. Check that
            # here.
            self.query_store(store_path, "--hash", hide_stderr=True)
            self._known_store_paths.add(store_path)
            return True
        except CalledProcessError:
            return False

    def query_store(self, store_path, query, hide_stderr=False):
        """Given a query (e.g. --hash or --size), perform the query.

        :param store_path: The store path to query.
        :type store_path: ``str``
        :param query: The query to perform. Must be a valid nix-store query.
        :type query: ``str``
        :param hide_stderr: If true, stderr will be hidden.<commit_msg>use a direct connection to the sqlite db, when available<commit_after>"""Serve nix store objects over HTTP."""
import argparse
import logging
import os
from os.path import exists, isdir, join, basename, dirname
import re
from subprocess import check_output, Popen, PIPE, CalledProcessError
from threading import Thread
# Special-case here to address a runtime bug I've encountered
try:
    import sqlite3
except ImportError as err:
    if "does not define init" in str(err):
        exit("Could not import sqlite3. This is probably due to PYTHONPATH "
             "corruption: make sure your PYTHONPATH is empty prior to "
             "running this command.")
    else:
        raise

from flask import Flask, make_response, send_file, request, jsonify
import six

from servenix import __version__
from servenix.common.utils import decode_str, strip_output, find_nix_paths
from servenix.common.exceptions import (NoSuchObject, NoNarGenerated,
                                        BaseHTTPError, NixImportFailed,
                                        CouldNotUpdateHash, ClientError)

_HASH_REGEX=re.compile(r"[a-z0-9]{32}")
_PATH_REGEX=re.compile(r"([a-z0-9]{32})-[^' \n/]*$")

class NixServer(Flask):
    """Serves nix packages."""
    def __init__(self, nix_store_path, nix_bin_path, nix_state_path,
                 compression_type, debug):
        # Path to the local nix store.
        self._nix_store_path = nix_store_path
        # Path to the local nix state directory.
        self._nix_state_path = nix_state_path
        # Path to the folder containing nix binaries.
        self._nix_bin_path = nix_bin_path
        # Matches valid nix store paths (local to this store)
        self._full_store_path_regex = re.compile(
            join(self._nix_store_path, _PATH_REGEX.pattern))
        self._compression_type = compression_type
        # Cache mapping object hashes to store paths.
        self._hashes_to_paths = {}
        # Cache mapping object hashes to store paths which have been validated.
        self._hashes_to_valid_paths = {}
        # Cache mapping store paths to object info.
        self._paths_to_info = {}
        # Set of known store paths.
        self._known_store_paths = set()
        # A static string telling a nix client what this store serves.
        self._cache_info = "\n".join([
            "StoreDir: {}".format(self._nix_store_path),
            "WantMassQuery: 1",
            "Priority: 30"
        ]) + "\n"
        if self._compression_type == "bzip2":
            self._content_type = "application/x-bzip2"
            self._nar_extension = ".nar.bz2"
        else:
            self._content_type = "application/x-xz"
            self._nar_extension = ".nar.xz"
        # Enable interactive debugging on unknown errors.
        self._debug = debug
        # Test connect to the nix database; if successful, then we
        # will use a direct connection to the database rather than
        # using nix-store. This is much faster, but is unavailable on
        # some systems.
        try:
            db_path = join(nix_state_path, "nix", "db", "db.sqlite")
            query = "select * from ValidPaths limit 1"
            db_con = sqlite3.connect(db_path)
            db_con.execute(query).fetchall()
            # If this succeeds, assign the db_con attribute.
            self._db_con = db_con
        except Exception as err:
            logging.warn("Couldn't connect to the database ({}). Can't "
                         "operate in direct-database mode :(".format(err))
            self._db_con = None

    def store_path_from_hash(self, store_object_hash):
        """Look up a store path using its hash.

        The name of every object in the nix store starts with a
        32-character hash. We can find the full path by finding an
        object that starts with this hash.

        :param store_object_hash: The 32-character hash prefix of the object.
        :type store_object_hash: ``str``

        :return: The full store path to the object.
        :rtype: ``str``

        :raises: :py:class:`NoSuchObject` if the object isn't in the store.
        """
        if store_object_hash in self._hashes_to_valid_paths:
            # Case: the hash maps to a path which has been validated
            # to exist in the nix store and in the database.
            return self._hashes_to_valid_paths[store_object_hash]
        elif store_object_hash in self._hashes_to_paths:
            # Case: the hash maps to a path which is known to exist in
            # the nix-store, but hasn't been checked for existence in
            # the database. Check that it's in the DB before returning.
            path = self._hashes_to_paths[store_object_hash]
            if self.check_in_store(path):
                self._hashes_to_valid_paths[store_object_hash] = path
                return path
            # Otherwise, remove it from the known hashes and raise an error.
            self._hashes_to_paths.pop(store_object_hash, None)
            raise NoSuchObject("No object with hash {} was found!"
                               .format(store_object_hash))
        if self._db_con is not None:
            # If we have a direct database connection, use this to check
            # path existence.
            query = ("select path from ValidPaths where path like '{}%'"
                     .format(join(self._nix_store_path, store_object_hash)))
            with self._db_con:
                paths = self._db_con.execute(query).fetchall()
            if len(paths) > 0:
                path = paths[0][0]
                self._hashes_to_valid_paths[store_object_hash] = path
                return path
        else:
            # Get the list of store objects by listing the directory.
            # Iterate through them until a matching hash is found, or
            # we've exhausted all paths, in which case we error.
            for path in map(decode_str, os.listdir(self._nix_store_path)):
                match = _PATH_REGEX.match(path)
                if match is None:
                    continue
                path = join(self._nix_store_path, path)
                prefix = match.group(1)
                # Add every path seen to the _hashes_to_paths cache.
                self._hashes_to_paths[prefix] = path
                if prefix == store_object_hash:
                    # The path exists in the store. Ensure it's also a
                    # valid path according to nix-store.
                    if not self.check_in_store(path):
                        break
                    self._hashes_to_valid_paths[store_object_hash] = path
                    return path
        # If we've gotten here, then the hash doesn't match any path.
        raise NoSuchObject("No object with hash {} was found."
                           .format(store_object_hash))

    def check_in_store(self, store_path):
        """Check that a store path exists in the nix store.

        :param store_path: Path to an object in the nix store. Is assumed
            to match the full path regex.
        :type store_path: ``str``

        :return: True if it's in the store, and False otherwise.
        :rtype: ``bool``
        """
        if store_path in self._known_store_paths:
            return True
        elif not exists(store_path):
            # If the path isn't in the filesystem, it definitely is
            # not a valid path.
            return False
        # If it is on the filesystem, it doesn't necessarily mean
        # that it's a registered path in the store. Check that
        # here.
        # If we have a connection to the database, all we have to
        # do is look in the database.
        if self._db_con is not None:
            query = ("select path from ValidPaths where path = '{}'"
                     .format(store_path))
            with self._db_con:
                results = self._db_con.execute(query).fetchall()
            if len(results) > 0:
                self._known_store_paths.add(store_path)
                return True
            else:
                return False
        else:
            # Otherwise we have to use the slower method :(
            try:
                self.query_store(store_path, "--hash", hide_stderr=True)
                self._known_store_paths.add(store_path)
                return True
            except CalledProcessError:
                return False

    def query_store(self, store_path, query, hide_stderr=False):
        """Given a query (e.g. --hash or --size), perform the query.

        :param store_path: The store path to query.
        :type store_path: ``str``
        :param query: The query to perform. Must be a valid nix-store query.
        :type query: ``str``
        :param hide_stderr: If true, stderr will be hidden.
{'fraction_non_alphanumeric': 0.07457067371202114, 'fraction_numerical': 0.003038309114927345, 'mean_word_length': 2.0227590337392694, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2467632', 'n_tokens_mistral': 4216, 'n_tokens_neox': 4057, 'n_words': 1428}