module Shebanger where

import Control.Exception (finally)
import Data.ByteString (ByteString)
import qualified Data.ByteString as ByteString
import Data.ByteString.Base64 (encode)
import Data.Foldable (for_)
import Data.List (isSuffixOf)
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NonEmpty
import Data.Maybe (fromMaybe)
import Shebanger.Cli (Command (..), ExecArgs (..), TranslateArgs (..), parseCliOpts)
import System.Directory (doesFileExist, removeFile)
import System.FilePath (takeFileName, (<.>), (-<.>))
import System.Posix (setFileMode, fileMode, getFileStatus, unionFileModes, ownerExecuteMode, groupExecuteMode, otherExecuteMode)
import System.Posix.Process (executeFile)
import System.Posix.ByteString (getEnv, setEnv, unsetEnv)
import System.Process (callProcess)
import Text.Read (readMaybe)

-- $setup
--
-- We need things from QuickCheck for some of the tests.
--
-- >>> import Test.QuickCheck

defaultMain :: IO ()
defaultMain :: IO ()
defaultMain = do
  Command
cmd <- IO Command
parseCliOpts
  Command -> IO ()
runCmd Command
cmd

runCmd :: Command -> IO ()
runCmd :: Command -> IO ()
runCmd = \case
  Translate TranslateArgs
transArgs -> TranslateArgs -> IO ()
runCmdTranslate TranslateArgs
transArgs
  Exec ExecArgs
execArgs -> ExecArgs -> IO ()
runCmdExec ExecArgs
execArgs

-- | The length of chunks of the input script.
--
-- The input script will be chunked into parts of this many bytes, then base-64
-- encoded (which increases the length by about 33%).  This base-64-encoded string
-- will then be put into the shebang line of all the shebanged scripts.
--
-- Note that Linux has a limit on how long a shebang line can be, so in practice this
-- has to be below 150 characters or so.
inputScriptChunkLength :: Int
inputScriptChunkLength :: Int
inputScriptChunkLength = Int
50

runCmdTranslate :: TranslateArgs -> IO ()
runCmdTranslate :: TranslateArgs -> IO ()
runCmdTranslate TranslateArgs
transArgs = do
  let inputScriptFileName :: [Char]
inputScriptFileName = [Char] -> [Char]
takeFileName TranslateArgs
transArgs.scriptFilePath
  ByteString
inputScriptContents <- [Char] -> IO ByteString
ByteString.readFile TranslateArgs
transArgs.scriptFilePath
  let chunkedInputScript :: [ByteString]
chunkedInputScript = Int -> ByteString -> [ByteString]
chunkByteString Int
inputScriptChunkLength ByteString
inputScriptContents
      b64ChunkedInputScript :: [ByteString]
b64ChunkedInputScript = (ByteString -> ByteString) -> [ByteString] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> ByteString
encode [ByteString]
chunkedInputScript
  [Char] -> [ByteString] -> IO ()
writeShebanged [Char]
inputScriptFileName [ByteString]
b64ChunkedInputScript

writeShebanged :: String -> [ByteString] -> IO ()
writeShebanged :: [Char] -> [ByteString] -> IO ()
writeShebanged [Char]
scriptName [ByteString]
b64Chunks =
  [(Int, ByteString)] -> ((Int, ByteString) -> IO ()) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ ([Int] -> [ByteString] -> [(Int, ByteString)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [ByteString]
b64Chunks) (((Int, ByteString) -> IO ()) -> IO ())
-> ((Int, ByteString) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Int
i :: Int, ByteString
chunk) -> do
    let fnameBase :: [Char]
fnameBase = [Char]
scriptName [Char] -> [Char] -> [Char]
<.> [Char]
"shebanged"
        fname :: [Char]
fname = if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then [Char]
fnameBase else [Char]
fnameBase [Char] -> [Char] -> [Char]
<.> Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i
        scriptContents :: ByteString
scriptContents = ByteString
"#!/usr/bin/env -S shebanger exec " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
chunk
    [Char] -> ByteString -> IO ()
ByteString.writeFile [Char]
fname ByteString
scriptContents
    [Char] -> IO ()
makeExecutable [Char]
fname

makeExecutable :: FilePath -> IO ()
makeExecutable :: [Char] -> IO ()
makeExecutable [Char]
path = do
    -- Get the current file permissions
    FileStatus
status <- [Char] -> IO FileStatus
getFileStatus [Char]
path
    let currentMode :: FileMode
currentMode = FileStatus -> FileMode
fileMode FileStatus
status
    -- Define the new mode with executable permissions
    let newMode :: FileMode
newMode = FileMode
currentMode FileMode -> FileMode -> FileMode
`unionFileModes` FileMode
ownerExecuteMode FileMode -> FileMode -> FileMode
`unionFileModes` FileMode
groupExecuteMode FileMode -> FileMode -> FileMode
`unionFileModes` FileMode
otherExecuteMode
    -- Set the file mode to the new mode
    [Char] -> FileMode -> IO ()
setFileMode [Char]
path FileMode
newMode

chunkByteString :: Int -> ByteString -> [ByteString]
chunkByteString :: Int -> ByteString -> [ByteString]
chunkByteString Int
n ByteString
bs
    | ByteString -> Bool
ByteString.null ByteString
bs = []
    | Bool
otherwise =
        let (ByteString
chunk, ByteString
rest) = Int -> ByteString -> (ByteString, ByteString)
ByteString.splitAt Int
n ByteString
bs
        in ByteString
chunk ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: Int -> ByteString -> [ByteString]
chunkByteString Int
n ByteString
rest

runCmdExec :: ExecArgs -> IO ()
runCmdExec :: ExecArgs -> IO ()
runCmdExec ExecArgs
execArgs = do
  Maybe ByteString
maybeShebangerScriptContents <- ByteString -> IO (Maybe ByteString)
getEnv ByteString
"SHEBANGER_SCRIPT_CONTENTS"
  case [Char] -> Either [Char] Int
getShebangedIndex ExecArgs
execArgs.shebangScriptFilePath of
    Left [Char]
badIndex ->
      [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$
        [Char]
"ERROR! Failed when trying to parse the current script's " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<>
        [Char]
"shebang index.  From file path \"" [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<>
        ExecArgs
execArgs.shebangScriptFilePath [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<>
        [Char]
"\", failed when trying to parse index: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<>
        [Char]
badIndex
    Right Int
idx -> do
      Maybe [Char]
maybeNextScript <- [Char] -> Int -> IO (Maybe [Char])
findNextScript ExecArgs
execArgs.shebangScriptFilePath Int
idx
      case Maybe [Char]
maybeNextScript of
        -- There is a next script that exists.  Update env var and execute the
        -- next script.
        Just [Char]
nextScript -> do
          case Maybe ByteString
maybeShebangerScriptContents of
            Maybe ByteString
Nothing ->
              if Int
idx Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
              then
                -- This is the initial script (script.sh.shebanger), so we
                -- expect that there is not yet a SHEBANGER_SCRIPT_CONTENTS
                -- env var.  We create the env var with the initial part of
                -- the script.
                ByteString -> ByteString -> Bool -> IO ()
setEnv ByteString
"SHEBANGER_SCRIPT_CONTENTS" ExecArgs
execArgs.shebangScriptPart Bool
True
              else
                -- This is not the initial script (so it has a name like
                -- script.sh.shebanger.07), but there is no
                -- SHEBANGER_SCRIPT_CONTENTS env var.  This is unexpected.
                [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$
                  [Char]
"ERROR! This is not the first script, but no " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<>
                  [Char]
"SHEBANGER_SCRIPT_CONTENTS env var was found."
            Just ByteString
envVarScriptContents ->
              ByteString -> ByteString -> Bool -> IO ()
setEnv
                ByteString
"SHEBANGER_SCRIPT_CONTENTS"
                (ByteString
envVarScriptContents ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ExecArgs
execArgs.shebangScriptPart)
                Bool
True
          -- exec the next script
          [Char] -> Bool -> [[Char]] -> Maybe [([Char], [Char])] -> IO ()
forall a.
[Char] -> Bool -> [[Char]] -> Maybe [([Char], [Char])] -> IO a
executeFile [Char]
nextScript Bool
True ExecArgs
execArgs.additionalArgs Maybe [([Char], [Char])]
forall a. Maybe a
Nothing
        -- There is not a next script.  This script is the last script.
        Maybe [Char]
Nothing -> do
          let fullScript :: ByteString
fullScript =
                ByteString -> Maybe ByteString -> ByteString
forall a. a -> Maybe a -> a
fromMaybe ByteString
"" Maybe ByteString
maybeShebangerScriptContents ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ExecArgs
execArgs.shebangScriptPart
              finalScriptName :: [Char]
finalScriptName =
                if Int
idx Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                then ExecArgs
execArgs.shebangScriptFilePath [Char] -> [Char] -> [Char]
<.> [Char]
"final"
                else ExecArgs
execArgs.shebangScriptFilePath [Char] -> [Char] -> [Char]
-<.> [Char]
"final"

          [Char] -> ByteString -> IO ()
ByteString.writeFile [Char]
finalScriptName ByteString
fullScript
          [Char] -> IO ()
makeExecutable [Char]
finalScriptName

          -- unset the SHEBANGER_SCRIPT_CONTENTS env var, since we don't want
          -- it to be inherited by the child.
          ByteString -> IO ()
unsetEnv ByteString
"SHEBANGER_SCRIPT_CONTENTS"

          -- Call the new executable, making sure to unlink it afterwards.
          IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
finally
            ([Char] -> [[Char]] -> IO ()
callProcess [Char]
finalScriptName ExecArgs
execArgs.additionalArgs)
            ([Char] -> IO ()
removeFile [Char]
finalScriptName)

findNextScript :: FilePath -> Int -> IO (Maybe FilePath)
findNextScript :: [Char] -> Int -> IO (Maybe [Char])
findNextScript [Char]
fp Int
currIdx = do
  let nextIdx :: Int
nextIdx = Int
currIdx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
      nextScriptName :: [Char]
nextScriptName =
        -- The initial script just ends with `.shebanged` (which we consider
        -- index 0), so in that case just add ".1".  Otherwise, increment the
        -- index by one.
        if Int
currIdx Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then [Char]
fp [Char] -> [Char] -> [Char]
<.> [Char]
".1"
        else [Char]
fp [Char] -> [Char] -> [Char]
-<.> Int -> [Char]
forall a. Show a => a -> [Char]
show Int
nextIdx
  Bool
exists <- [Char] -> IO Bool
doesFileExist [Char]
nextScriptName
  Maybe [Char] -> IO (Maybe [Char])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe [Char] -> IO (Maybe [Char]))
-> Maybe [Char] -> IO (Maybe [Char])
forall a b. (a -> b) -> a -> b
$ if Bool
exists then [Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
nextScriptName else Maybe [Char]
forall a. Maybe a
Nothing

-- | Returns 'Left' 'String' of the last part of the filename that it is trying
-- to parse a number on error.
--
-- Returns 'Right' 'Int' with the index if it is found correctly.
getShebangedIndex :: FilePath -> Either String Int
getShebangedIndex :: [Char] -> Either [Char] Int
getShebangedIndex [Char]
fname =
  if [Char]
".shebanged" [Char] -> [Char] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isSuffixOf` [Char]
fname
  -- The .shebanged file is the initial file that represents index 0
  then Int -> Either [Char] Int
forall a b. b -> Either a b
Right Int
0
  else do
    -- The shebanged filenames are going to look like:
    -- my-script.sh.shebanged.XX, where XX is a number.
    let splitFname :: NonEmpty [Char]
splitFname = (Char -> Bool) -> [Char] -> NonEmpty [Char]
forall a. (a -> Bool) -> [a] -> NonEmpty [a]
split (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'.') [Char]
fname
        -- pull out just the number
        strNum :: [Char]
strNum = NonEmpty [Char] -> [Char]
forall a. NonEmpty a -> a
NonEmpty.last NonEmpty [Char]
splitFname
        -- try to parse it as a number
        maybeNum :: Maybe Int
maybeNum = [Char] -> Maybe Int
forall a. Read a => [Char] -> Maybe a
readMaybe [Char]
strNum
    case Maybe Int
maybeNum of
      -- could not read file name part as number
      Maybe Int
Nothing -> [Char] -> Either [Char] Int
forall a b. a -> Either a b
Left [Char]
strNum
      Just Int
num
        -- the file name part should never be less than 1
        | Int
num Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 -> [Char] -> Either [Char] Int
forall a b. a -> Either a b
Left [Char]
strNum
        | Bool
otherwise -> Int -> Either [Char] Int
forall a b. b -> Either a b
Right Int
num

-- | Split a list based on a predicate.
--
-- >>> split (== ' ') "hello world my name is bob"
-- "hello" :| ["world","my","name","is","bob"]
--
-- A predicate of @'const' 'True'@ splits on everything, leaving you with
-- a list of empty lists, with one more entry than your original list:
--
-- >>> split (const True) "bye"
-- "" :| ["","",""]
--
-- A predicate of @'const' 'False'@ produces no splits:
--
-- >>> split (const False) "bye"
-- "bye" :| []
--
-- An empty list doesn't get split, regardless of the predicate:
--
-- prop> \(Fun _ f) -> split f "" == ("" :| [])
split :: forall a. (a -> Bool) -> [a] -> NonEmpty [a]
split :: forall a. (a -> Bool) -> [a] -> NonEmpty [a]
split a -> Bool
_ [] = [] [a] -> [[a]] -> NonEmpty [a]
forall a. a -> [a] -> NonEmpty a
:| []
split a -> Bool
p [a]
t = [a] -> NonEmpty [a]
loop [a]
t
  where
    loop :: [a] -> NonEmpty [a]
    loop :: [a] -> NonEmpty [a]
loop [a]
s
      | [a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
s'   = [a]
l [a] -> [[a]] -> NonEmpty [a]
forall a. a -> [a] -> NonEmpty a
:| []
      | Bool
otherwise = [a] -> NonEmpty [a] -> NonEmpty [a]
forall a. a -> NonEmpty a -> NonEmpty a
NonEmpty.cons [a]
l (NonEmpty [a] -> NonEmpty [a]) -> NonEmpty [a] -> NonEmpty [a]
forall a b. (a -> b) -> a -> b
$ [a] -> NonEmpty [a]
loop ([a] -> [a]
forall a. HasCallStack => [a] -> [a]
tail [a]
s')
      where ([a]
l, [a]
s') = (a -> Bool) -> [a] -> ([a], [a])
forall a. (a -> Bool) -> [a] -> ([a], [a])
break a -> Bool
p [a]
s