text: # Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
  # Licensed under the Apache License, Version 2.0 (see LICENSE).
  
  from __future__ import annotations
  
  import dataclasses
  import hashlib
  import os
  from dataclasses import dataclass
  from enum import Enum
  from textwrap import dedent
  from typing import Iterable, Sequence
  
  from pants.core.subsystems import python_bootstrap
  from pants.core.subsystems.python_bootstrap import PythonBootstrap
  from pants.engine.collection import DeduplicatedCollection
  from pants.engine.engine_aware import EngineAwareReturnType
  from pants.engine.fs import CreateDigest, FileContent
  from pants.engine.internals.native_engine import Digest
  from pants.engine.internals.selectors import Get, MultiGet
  from pants.engine.process import FallibleProcessResult, Process, ProcessCacheScope, ProcessResult
  from pants.engine.rules import collect_rules, rule
  from pants.util.logging import LogLevel
  from pants.util.meta import frozen_after_init
  from pants.util.ordered_set import OrderedSet
  from pants.util.strutil import create_path_env_var, pluralize
  
  # -------------------------------------------------------------------------------------------
  # `BinaryPath` types
  # -------------------------------------------------------------------------------------------
  
  # TODO(#14492): This should be configurable via `[system-binaries]` subsystem, likely per-binary.
  SEARCH_PATHS = ("/usr/bin", "/bin", "/usr/local/bin")
  
  
  @frozen_after_init
  @dataclass(unsafe_hash=True)
  class BinaryPath:
      path: str
      fingerprint: str
  
      def __init__(self, path: str, fingerprint: str | None = None) -> None:
          self.path = path
          self.fingerprint = self._fingerprint() if fingerprint is None else fingerprint
  
      @staticmethod
      def _fingerprint(content: bytes | bytearray | memoryview | None = None) -> str:
          hasher = hashlib.sha256() if content is None else hashlib.sha256(content)
          return hasher.hexdigest()
  
      @classmethod
      def fingerprinted(
          cls, path: str, representative_content: bytes | bytearray | memoryview
      ) -> BinaryPath:
          return cls(path, fingerprint=cls._fingerprint(representative_content))
  
  
  @frozen_after_init
  @dataclass(unsafe_hash=True)
  class BinaryPathTest:
      args: tuple[str, ...]
      fingerprint_stdout: bool
  
      def __init__(self, args: Iterable[str], fingerprint_stdout: bool = True) -> None:
          self.args = tuple(args)
          self.fingerprint_stdout = fingerprint_stdout
  
  
  class SearchPath(DeduplicatedCollection[str]):
      """The search path for binaries; i.e.: the $PATH."""
  
  
  @frozen_after_init
  @dataclass(unsafe_hash=True)
  class BinaryPathRequest:
      """Request to find a binary of a given name.
  
      If `check_file_entries` is `True` a BinaryPathRequest will consider any entries in the
      `search_path` that are file paths in addition to traditional directory paths.
  
      If a `test` is specified all binaries that are found will be executed with the test args and
      only those binaries whose test executions exit with return code 0 will be retained.
      Additionally, if test execution includes stdout content, that will be used to fingerprint the
      binary path so that upgrades and downgrades can be detected. A reasonable test for many programs
      might be `BinaryPathTest(args=["--version"])` since it will both ensure the program runs and
      also produce stdout text that changes upon upgrade or downgrade of the binary at the discovered
      path.
      """
  
      search_path: SearchPath
      binary_name: str
      check_file_entries: bool
      test: BinaryPathTest | None
  
      def __init__(
          self,
          *,
          search_path: Iterable[str],
          binary_name: str,
          check_file_entries: bool = False,
          test: BinaryPathTest | None = None,
      ) -> None:
          self.search_path = SearchPath(search_path)
          self.binary_name = binary_name
          self.check_file_entries = check_file_entries
          self.test = test
  
  
  @frozen_after_init
  @dataclass(unsafe_hash=True)
  class BinaryPaths(EngineAwareReturnType):
      binary_name: str
      paths: tuple[BinaryPath, ...]
  
      def __init__(self, binary_name: str, paths: Iterable[BinaryPath] | None = None):
          self.binary_name = binary_name
          self.paths = tuple(OrderedSet(paths) if paths else ())
  
      def message(self) -> str:
          if not self.paths:
              return f"failed to find {self.binary_name}"
          found_msg = f"found {self.binary_name} at {self.paths[0]}"
          if len(self.paths) > 1:
              found_msg = f"{found_msg} and {pluralize(len(self.paths) - 1, 'other location')}"
          return found_msg
  
      @property
      def first_path(self) -> BinaryPath | None:
          """Return the first path to the binary that was discovered, if any."""
          return next(iter(self.paths), None)
  
      def first_path_or_raise(self, request: BinaryPathRequest, *, rationale: str) -> BinaryPath:
          """Return the first path to the binary that was discovered, if any."""
          first_path = self.first_path
          if not first_path:
              raise BinaryNotFoundError.from_request(request, rationale=rationale)
          return first_path
  
  
  class BinaryNotFoundError(EnvironmentError):
      @classmethod
      def from_request(
          cls,
          request: BinaryPathRequest,
          *,
          rationale: str | None = None,
          alternative_solution: str | None = None,
      ) -> BinaryNotFoundError:
          """When no binary is found via `BinaryPaths`, and it is not recoverable.
  
          :param rationale: A short description of why this binary is needed, e.g.
              "download the tools Pants needs" or "run Python programs".
          :param alternative_solution: A description of what else users can do to fix the issue,
              beyond installing the program. For example, "Alternatively, you can set the option
              `--python-bootstrap-search-path` to change the paths searched."
          """
          msg = (
              f"Cannot find `{request.binary_name}` on `{sorted(request.search_path)}`. Please "
              "ensure that it is installed"
          )
          msg += f" so that Pants can {rationale}." if rationale else "."
          if alternative_solution:
              msg += f"\n\n{alternative_solution}"
          return BinaryNotFoundError(msg)
  
  
  # -------------------------------------------------------------------------------------------
  # Binaries
  # -------------------------------------------------------------------------------------------
  
  
  class BashBinary(BinaryPath):
      """The `bash` binary."""
  
      DEFAULT_SEARCH_PATH = SearchPath(("/usr/bin", "/bin", "/usr/local/bin"))
  
  
  @dataclass(frozen=True)
  class BashBinaryRequest:
      search_path: SearchPath = BashBinary.DEFAULT_SEARCH_PATH
  
  
  class PythonBinary(BinaryPath):
      """A Python3 interpreter for use by `@rule` code as an alternative to BashBinary scripts.
  
      Python is usable for `@rule` scripting independently of `pants.backend.python`, but currently
      thirdparty dependencies are not supported, because PEX lives in that backend.
  
      TODO: Consider extracting PEX out into the core in order to support thirdparty dependencies.
      """
  
  
  # Note that updating this will impact the `archive` target defined in `core/target_types.py`.
  class ArchiveFormat(Enum):
      TAR = "tar"
      TGZ = "tar.gz"
      TBZ2 = "tar.bz2"
      TXZ = "tar.xz"
      ZIP = "zip"
  
  
  class ZipBinary(BinaryPath):
      def create_archive_argv(
          self, output_filename: str, input_files: Sequence[str]
      ) -> tuple[str, ...]:
          return (self.path, output_filename, *input_files)
  
  
  class UnzipBinary(BinaryPath):
      def extract_archive_argv(self, archive_path: str, extract_path: str) -> tuple[str, ...]:
          # Note that the `output_dir` does not need to already exist.
          # The caller should validate that it's a valid `.zip` file.
          return (self.path, archive_path, "-d", extract_path)
  
  
  @dataclass(frozen=True)
  class GunzipBinary:
      python: PythonBinary
  
      def extract_archive_argv(self, archive_path: str, extract_path: str) -> tuple[str, ...]:
          archive_name = os.path.basename(archive_path)
          dest_file_name = os.path.splitext(archive_name)[0]
          dest_path = os.path.join(extract_path, dest_file_name)
          script = dedent(
              f"""
              import gzip
              import shutil
              with gzip.GzipFile(filename={archive_path!r}, mode="rb") as source:
                  with open({dest_path!r}, "wb") as dest:
                      shutil.copyfileobj(source, dest)
              """
          )
          return (self.python.path, "-c", script)
  
  
  class TarBinary(BinaryPath):
      def create_archive_argv(
          self, output_filename: str, input_files: Sequence[str], tar_format: ArchiveFormat
      ) -> tuple[str, ...]:
          # Note that the parent directory for the output_filename must already exist.
          #
          # We do not use `-a` (auto-set compression) because it does not work with older tar
          # versions. Not all tar implementations will support these compression formats - in that
          # case, the user will need to choose a different format.
          compression = {ArchiveFormat.TGZ: "z", ArchiveFormat.TBZ2: "j", ArchiveFormat.TXZ: "J"}.get(
              tar_format, ""
          )
          return (self.path, f"c{compression}f", output_filename, *input_files)
  
      def extract_archive_argv(self, archive_path: str, extract_path: str) -> tuple[str, ...]:
          # Note that the `output_dir` must already exist.
          # The caller should validate that it's a valid `.tar` file.
          return (self.path, "xf", archive_path, "-C", extract_path)
  
  
  # -------------------------------------------------------------------------------------------
  # Rules to find binaries
  # -------------------------------------------------------------------------------------------
  
  
  @rule(desc="Finding the `bash` binary", level=LogLevel.DEBUG)
  async def find_bash(bash_request: BashBinaryRequest) -> BashBinary:
      request = BinaryPathRequest(
          binary_name="bash",
          search_path=bash_request.search_path,
          test=BinaryPathTest(args=["--version"]),
      )
      paths = await Get(BinaryPaths, BinaryPathRequest, request)
      first_path = paths.first_path
      if not first_path:
          raise BinaryNotFoundError.from_request(request)
      return BashBinary(first_path.path, first_path.fingerprint)
  
  
  @rule
  async def get_bash() -> BashBinary:
      # Expose bash to external consumers.
      return await Get(BashBinary, BashBinaryRequest())
  
  
  @rule
  async def find_binary(request: BinaryPathRequest) -> BinaryPaths:
      # If we are not already locating bash, recurse to locate bash to use it as an absolute path in
      # our shebang. This avoids mixing locations that we would search for bash into the search paths
      # of the request we are servicing.
      # TODO(#10769): Replace this script with a statically linked native binary so we don't
      #  depend on either /bin/bash being available on the Process host.
      if request.binary_name == "bash":
          shebang = "#!/usr/bin/env bash"
      else:
          bash = await Get(BashBinary, BashBinaryRequest())
          shebang = f"#!{bash.path}"
  
      script_path = "./find_binary.sh"
      script_header = dedent(
          f"""\
          {shebang}
  
          set -euox pipefail
  
          CHECK_FILE_ENTRIES={'1' if request.check_file_entries else ''}
          """
      )
      script_body = dedent(
          """\
          for path in ${PATH//:/ }; do
              if [[ -d "${path}" ]]; then
                  # Handle traditional directory PATH element.
                  maybe_exe="${path}/$1"
              elif [[ -n "${CHECK_FILE_ENTRIES}" ]]; then
                  # Handle PATH elements that are filenames to allow for precise selection.
                  maybe_exe="${path}"
              else
                  maybe_exe=
              fi
              if [[ "$1" == "${maybe_exe##*/}" && -f "${maybe_exe}" && -x "${maybe_exe}" ]]
              then
                  echo "${maybe_exe}"
              fi
          done
          """
      )
      script_content = script_header + script_body
      script_digest = await Get(
          Digest,
          CreateDigest([FileContent(script_path, script_content.encode(), is_executable=True)]),
      )
  
      # Some subtle notes about executing this script:
      #
      #  - We run the script with `ProcessResult` instead of `FallibleProcessResult` so that we
      #      can catch bugs in the script itself, given an earlier silent failure.
      #  - We set `ProcessCacheScope.PER_RESTART_SUCCESSFUL` to force re-run since any binary found
      #      on the host system today could be gone tomorrow. Ideally we'd only do this for local
      #      processes since all known remoting configurations include a static container image as
      #      part of their cache key which automatically avoids this problem. See #10769 for a
      #      solution that is less of a tradeoff.
      search_path = create_path_env_var(request.search_path)
      result = await Get(
          ProcessResult,
          Process(
              description=f"Searching for `{request.binary_name}` on PATH={search_path}",
              level=LogLevel.DEBUG,
              input_digest=script_digest,
              argv=[script_path, request.binary_name],
              env={"PATH": search_path},
              cache_scope=ProcessCacheScope.PER_RESTART_SUCCESSFUL,
          ),
      )
  
      binary_paths = BinaryPaths(binary_name=request.binary_name)
      found_paths = result.stdout.decode().splitlines()
      if not request.test:
          return dataclasses.replace(binary_paths, paths=[BinaryPath(path) for path in found_paths])
  
      results = await MultiGet(
          Get(
              FallibleProcessResult,
              Process(
                  description=f"Test binary {path}.",
                  level=LogLevel.DEBUG,
                  argv=[path, *request.test.args],
                  # NB: Since a failure is a valid result for this script, we always cache it for
                  # `pantsd`'s lifetime, regardless of success or failure.
                  cache_scope=ProcessCacheScope.PER_RESTART_ALWAYS,
              ),
          )
          for path in found_paths
      )
      return dataclasses.replace(
          binary_paths,
          paths=[
              (
                  BinaryPath.fingerprinted(path, result.stdout)
                  if request.test.fingerprint_stdout
                  else BinaryPath(path, result.stdout.decode())
              )
              for path, result in zip(found_paths, results)
              if result.exit_code == 0
          ],
      )
  
  
  @rule(desc="Finding a `python` binary", level=LogLevel.TRACE)
  async def find_python(python_bootstrap: PythonBootstrap) -> PythonBinary:
      # PEX files are compatible with bootstrapping via Python 2.7 or Python 3.5+, but we select 3.6+
      # for maximum compatibility with internal scripts.
      interpreter_search_paths = python_bootstrap.interpreter_search_paths()
      all_python_binary_paths = await MultiGet(
          Get(
              BinaryPaths,
              BinaryPathRequest(
                  search_path=interpreter_search_paths,
                  binary_name=binary_name,
                  check_file_entries=True,
                  test=BinaryPathTest(
                      args=[
                          "-c",
                          # N.B.: The following code snippet must be compatible with Python 3.6+.
                          #
                          # We hash the underlying Python interpreter executable to ensure we detect
                          # changes in the real interpreter that might otherwise be masked by Pyenv
                          # shim scripts found on the search path. Naively, just printing out the full
                          # version_info would be enough, but that does not account for supported abi
                          # changes (e.g.: a pyenv switch from a py27mu interpreter to a py27m
                          # interpreter.)
                          #
                          # When hashing, we pick 8192 for efficiency of reads and fingerprint updates
                          # (writes) since it's a common OS buffer size and an even multiple of the
                          # hash block size.
                          dedent(
                              """\
                              import sys
  
                              major, minor = sys.version_info[:2]
                              if not (major == 3 and minor >= 6):
                                  sys.exit(1)
  
                              import hashlib
                              hasher = hashlib.sha256()
                              with open(sys.executable, "rb") as fp:
                                  for chunk in iter(lambda: fp.read(8192), b""):
                                      hasher.update(chunk)
                              sys.stdout.write(hasher.hexdigest())
                              """
                          ),
                      ],
                      fingerprint_stdout=False,  # We already emit a usable fingerprint to stdout.
                  ),
              ),
          )
          for binary_name in python_bootstrap.interpreter_names
      )
  
      for binary_paths in all_python_binary_paths:
          path = binary_paths.first_path
          if path:
              return PythonBinary(
                  path=path.path,
                  fingerprint=path.fingerprint,
              )
  
      raise BinaryNotFoundError(
          "Was not able to locate a Python interpreter to execute rule code.\n"
          "Please ensure that Python is available in one of the locations identified by "
          "`[python-bootstrap] search_path`, which currently expands to:\n"
          f"  {interpreter_search_paths}"
      )
  
  
  @rule(desc="Finding the `zip` binary", level=LogLevel.DEBUG)
  async def find_zip() -> ZipBinary:
      request = BinaryPathRequest(
          binary_name="zip", search_path=SEARCH_PATHS, test=BinaryPathTest(args=["-v"])
      )
      paths = await Get(BinaryPaths, BinaryPathRequest, request)
      first_path = paths.first_path_or_raise(request, rationale="create `.zip` archives")
      return ZipBinary(first_path.path, first_path.fingerprint)
  
  
  @rule(desc="Finding the `unzip` binary", level=LogLevel.DEBUG)
  async def find_unzip() -> UnzipBinary:
      request = BinaryPathRequest(
          binary_name="unzip", search_path=SEARCH_PATHS, test=BinaryPathTest(args=["-v"])
      )
      paths = await Get(BinaryPaths, BinaryPathRequest, request)
      first_path = paths.first_path_or_raise(
          request, rationale="download the tools Pants needs to run"
      )
      return UnzipBinary(first_path.path, first_path.fingerprint)
  
  
  @rule
  def find_gunzip(python: PythonBinary) -> GunzipBinary:
      return GunzipBinary(python)
  
  
  @rule(desc="Finding the `tar` binary", level=LogLevel.DEBUG)
  async def find_tar() -> TarBinary:
      request = BinaryPathRequest(
          binary_name="tar", search_path=SEARCH_PATHS, test=BinaryPathTest(args=["--version"])
      )
      paths = await Get(BinaryPaths, BinaryPathRequest, request)
      first_path = paths.first_path_or_raise(
          request, rationale="download the tools Pants needs to run"
      )
      return TarBinary(first_path.path, first_path.fingerprint)
  
  
  # -------------------------------------------------------------------------------------------
  # Rules for lazy requests
  # TODO(#12946): Get rid of this when it becomes possible to use `Get()` with only one arg.
  # -------------------------------------------------------------------------------------------
  
  
  class ZipBinaryRequest:
      pass
  
  
  class UnzipBinaryRequest:
      pass
  
  
  class GunzipBinaryRequest:
      pass
  
  
  class TarBinaryRequest:
      pass
  
  
  @rule
  async def find_zip_wrapper(_: ZipBinaryRequest, zip_binary: ZipBinary) -> ZipBinary:
      return zip_binary
  
  
  @rule
  async def find_unzip_wrapper(_: UnzipBinaryRequest, unzip_binary: UnzipBinary) -> UnzipBinary:
      return unzip_binary
  
  
  @rule
  async def find_gunzip_wrapper(_: GunzipBinaryRequest, gunzip: GunzipBinary) -> GunzipBinary:
      return gunzip
  
  
  @rule
  async def find_tar_wrapper(_: TarBinaryRequest, tar_binary: TarBinary) -> TarBinary:
      return tar_binary
  
  
  def rules():
      return [*collect_rules(), *python_bootstrap.rules()]
  