text: # Copyright 2013-2021 Lawrence Livermore National Security, LLC and other
  # Spack Project Developers. See the top-level COPYRIGHT file for details.
  #
  # SPDX-License-Identifier: (Apache-2.0 OR MIT)
  
  import filecmp
  import os
  import re
  import shutil
  import stat
  import sys
  import tempfile
  
  import llnl.util.filesystem as fs
  import llnl.util.tty as tty
  
  import spack.paths
  import spack.store
  
  #: OS-imposed character limit for shebang line: 127 for Linux; 511 for Mac.
  #: Different Linux distributions have different limits, but 127 is the
  #: smallest among all modern versions.
  if sys.platform == 'darwin':
      system_shebang_limit = 511
  else:
      system_shebang_limit = 127
  
  #: Spack itself also limits the shebang line to at most 4KB, which should be plenty.
  spack_shebang_limit = 4096
  
  interpreter_regex = re.compile(b'#![ \t]*?([^ \t\0\n]+)')
  
  
  def sbang_install_path():
      """Location sbang should be installed within Spack's ``install_tree``."""
      sbang_root = str(spack.store.unpadded_root)
      install_path = os.path.join(sbang_root, "bin", "sbang")
      path_length = len(install_path)
      if path_length > system_shebang_limit:
          msg = ('Install tree root is too long. Spack cannot patch shebang lines'
                 ' when script path length ({0}) exceeds limit ({1}).\n  {2}')
          msg = msg.format(path_length, system_shebang_limit, install_path)
          raise SbangPathError(msg)
      return install_path
  
  
  def sbang_shebang_line():
      """Full shebang line that should be prepended to files to use sbang.
  
      The line returned does not have a final newline (caller should add it
      if needed).
  
      This should be the only place in Spack that knows about what
      interpreter we use for ``sbang``.
      """
      return '#!/bin/sh %s' % sbang_install_path()
  
  
  def get_interpreter(binary_string):
      # The interpreter may be preceded with ' ' and \t, is itself any byte that
      # follows until the first occurrence of ' ', \t, \0, \n or end of file.
      match = interpreter_regex.match(binary_string)
      return None if match is None else match.group(1)
  
  
  def filter_shebang(path):
      """
      Adds a second shebang line, using sbang, at the beginning of a file, if necessary.
      Note: Spack imposes a relaxed shebang line limit, meaning that a newline or end of
      file must occur before ``spack_shebang_limit`` bytes. If not, the file is not
      patched.
      """
      with open(path, 'rb') as original:
          # If there is no shebang, we shouldn't replace anything.
          old_shebang_line = original.read(2)
          if old_shebang_line != b'#!':
              return False
  
          # Stop reading after b'\n'. Note that old_shebang_line includes the first b'\n'.
          old_shebang_line += original.readline(spack_shebang_limit - 2)
  
          # If the shebang line is short, we don't have to do anything.
          if len(old_shebang_line) <= system_shebang_limit:
              return False
  
          # Whenever we can't find a newline within the maximum number of bytes, we will
          # not attempt to rewrite it. In principle we could still get the interpreter if
          # only the arguments are truncated, but note that for PHP we need the full line
          # since we have to append `?>` to it. Since our shebang limit is already very
          # generous, it's unlikely to happen, and it should be fine to ignore.
          if (
              len(old_shebang_line) == spack_shebang_limit and
              old_shebang_line[-1] != b'\n'
          ):
              return False
  
          # This line will be prepended to file
          new_sbang_line = (sbang_shebang_line() + '\n').encode('utf-8')
  
          # Skip files that are already using sbang.
          if old_shebang_line == new_sbang_line:
              return
  
          interpreter = get_interpreter(old_shebang_line)
  
          # If there was only whitespace we don't have to do anything.
          if not interpreter:
              return False
  
          # Store the file permissions, the patched version needs the same.
          saved_mode = os.stat(path).st_mode
  
          # Change non-writable files to be writable if needed.
          if not os.access(path, os.W_OK):
              os.chmod(path, saved_mode | stat.S_IWUSR)
  
          # No need to delete since we'll move it and overwrite the original.
          patched = tempfile.NamedTemporaryFile('wb', delete=False)
          patched.write(new_sbang_line)
  
          # Note that in Python this does not go out of bounds even if interpreter is a
          # short byte array.
          # Note: if the interpreter string was encoded with UTF-16, there would have
          # been a \0 byte between all characters of lua, node, php; meaning that it would
          # lead to truncation of the interpreter. So we don't have to worry about weird
          # encodings here, and just looking at bytes is justified.
          if interpreter[-4:] == b'/lua' or interpreter[-7:] == b'/luajit':
              # Use --! instead of #! on second line for lua.
              patched.write(b'--!' + old_shebang_line[2:])
          elif interpreter[-5:] == b'/node':
              # Use //! instead of #! on second line for node.js.
              patched.write(b'//!' + old_shebang_line[2:])
          elif interpreter[-4:] == b'/php':
              # Use <?php #!... ?> instead of #!... on second line for php.
              patched.write(b'<?php ' + old_shebang_line + b' ?>')
          else:
              patched.write(old_shebang_line)
  
          # After copying the remainder of the file, we can close the original
          shutil.copyfileobj(original, patched)
  
      # And close the temporary file so we can move it.
      patched.close()
  
      # Overwrite original file with patched file, and keep the original mode
      shutil.move(patched.name, path)
      os.chmod(path, saved_mode)
      return True
  
  
  def filter_shebangs_in_directory(directory, filenames=None):
      if filenames is None:
          filenames = os.listdir(directory)
      for file in filenames:
          path = os.path.join(directory, file)
  
          # only handle files
          if not os.path.isfile(path):
              continue
  
          # only handle executable files
          st = os.stat(path)
          if not st.st_mode & (stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH):
              continue
  
          # only handle links that resolve within THIS package's prefix.
          if os.path.islink(path):
              real_path = os.path.realpath(path)
              if not real_path.startswith(directory + os.sep):
                  continue
  
          # test the file for a long shebang, and filter
          if filter_shebang(path):
              tty.debug("Patched overlong shebang in %s" % path)
  
  
  def install_sbang():
      """Ensure that ``sbang`` is installed in the root of Spack's install_tree.
  
      This is the shortest known publicly accessible path, and installing
      ``sbang`` here ensures that users can access the script and that
      ``sbang`` itself is in a short path.
      """
      # copy in a new version of sbang if it differs from what's in spack
      sbang_path = sbang_install_path()
      if os.path.exists(sbang_path) and filecmp.cmp(
              spack.paths.sbang_script, sbang_path):
          return
  
      # make $install_tree/bin and copy in a new version of sbang if needed
      sbang_bin_dir = os.path.dirname(sbang_path)
      fs.mkdirp(sbang_bin_dir)
      fs.install(spack.paths.sbang_script, sbang_path)
      fs.set_install_permissions(sbang_bin_dir)
  
  
  def post_install(spec):
      """This hook edits scripts so that they call /bin/bash
      $spack_prefix/bin/sbang instead of something longer than the
      shebang limit.
      """
      if spec.external:
          tty.debug('SKIP: shebang filtering [external package]')
          return
  
      install_sbang()
  
      for directory, _, filenames in os.walk(spec.prefix):
          filter_shebangs_in_directory(directory, filenames)
  
  
  class SbangPathError(spack.error.SpackError):
      """Raised when the install tree root is too long for sbang to work."""
  