content: {-# LANGUAGE LambdaCase         #-}
  {-# LANGUAGE OverloadedStrings  #-}
  {-# LANGUAGE QuasiQuotes        #-}
  {-# LANGUAGE RecordWildCards    #-}
  
  {-| This module provides a basic library API to @nix-delegate@'s functionality
  -}
  
  module Nix.Delegate
      ( -- * Options
        OptArgs(..)
      , Command(..)
      , OperatingSystem(..)
  
        -- * Commands
      , delegate
      , delegateStream
      , main
      ) where
  
  import           Control.Applicative       (empty, many, (<**>), (<|>))
  import           Control.Exception         (SomeException)
  import           Control.Monad
  import           Control.Monad.IO.Class    (MonadIO)
  import           Control.Monad.Managed     (MonadManaged)
  import qualified Data.Foldable             as Foldable
  import           Data.Maybe
  import           Data.Monoid               ((<>))
  import           Data.Text                 (Text)
  import           Prelude                   hiding (FilePath)
  import           Turtle                    (ExitCode (..), FilePath, Line,
                                              Shell, d, fp, liftIO, s, (%), (</>))
  
  import qualified Control.Exception
  import qualified Control.Foldl             as Foldl
  import qualified Data.ByteString.Lazy
  import qualified Data.Text
  import qualified NeatInterpolation
  import qualified Options.Applicative       as Options
  import qualified Options.Applicative.Types as Options
  import qualified Turtle
  import qualified Turtle.Line
  
  -- | @delegate@ options
  data OptArgs = OptArgs
      { host    :: Text
      -- ^ Build host to add
      , os      :: [OperatingSystem]
      -- ^ Supported platform types (Default: @x86_64-linux@)
      , key     :: Maybe FilePath
      -- ^ SSH private key used to log in to build host (Default: @~/.ssh/id_rsa@)
      , cores   :: Maybe Integer
      -- ^ Number of cores available on the build host (Default: @1@)
      , feature :: [Text]
      -- ^ Supported system features for the build host
      , cmd     :: Command
      -- ^ Command to run with distributed builds enabled
      } deriving (Show)
  
  -- | Operating system
  data OperatingSystem
      = X86_64_Linux
      | X86_64_Darwin
      deriving (Show)
  
  parseOptions :: Options.Parser OptArgs
  parseOptions =
    OptArgs
    <$> parseHost
    <*> many parseOS
    <*> parseKey
    <*> parseCores
    <*> many parseFeature
    <*> parseCommand
  
  parseHost :: Options.Parser Text
  parseHost = Data.Text.pack <$>
    (Options.option Options.str $
     ( Options.long "host"
    <> Options.help "Machine to use as a build slave"
     )
    )
  
  parseKey :: Options.Parser (Maybe FilePath)
  parseKey =
    (Options.optional $  (Turtle.fromText . Data.Text.pack) <$>
     (Options.option Options.str $
      ( Options.long "key"
     <> Options.help "Path to SSH private key (Default: ~/.ssh/id_rsa)"
      )
     )
    )
  
  parseCores :: Options.Parser (Maybe Integer)
  parseCores =
    (Options.optional
     (Options.option ( Options.auto) $
      ( Options.long "cores"
     <> Options.help "Number of cores to use (Default: 1)"
      )
     )
    )
  
  parseFeature :: Options.Parser Text
  parseFeature = Data.Text.pack <$>
    (Options.option Options.str $
     ( Options.long "feature"
    <> Options.help "Supported system features"
     )
    )
  
  parseOS :: Options.Parser OperatingSystem
  parseOS =
        Options.flag' X86_64_Linux  (Options.long "x86_64-linux" )
    <|> Options.flag' X86_64_Darwin (Options.long "x86_64-darwin")
  
  renderOS :: OperatingSystem -> Text
  renderOS X86_64_Linux  = "x86_64-linux"
  renderOS X86_64_Darwin = "x86_64-darwin"
  
  -- | Command to run
  data Command = Command Text [Text]
    deriving (Show)
  
  parseCommand :: Options.Parser Command
  parseCommand = parseCmd <*> many (Data.Text.pack <$> Options.strArgument (Options.metavar "ARGS"))
    where
      cmdP :: Options.ReadM ([Text] -> Command)
      cmdP = Command . Data.Text.pack <$> Options.readerAsk
  
      parseCmd =
        (Options.argument cmdP $
          ( Options.metavar "COMMAND"
         <> Options.help    "Command to delegate (if 'nix-build', sudo will be used if $NIX_REMOTE=daemon)"
          )
        )
  
  renderCmd :: Command -> Text
  renderCmd (Command cmd args) = Turtle.format (s%" "%s) cmd (Data.Text.intercalate " " args)
  
  canSudo :: Command -> Bool
  canSudo (Command command _) = Turtle.filename path == "nix-build"
    where
      path = Turtle.fromText command
  
  -- | @main@ used by the @delegate@ executable
  main :: IO ()
  main = do
      options <- do
          Options.execParser
              (Options.info (parseOptions <**> Options.helper)
                  (   Options.fullDesc
                  <>  Options.progDesc "Run a subcommand with distributed builds transiently enabled"
                  <>  Options.noIntersperse
                  )
              )
      delegate options
  
  exchangeKeys :: FilePath -> Text -> IO ()
  exchangeKeys key host = do
    let key' = Turtle.format fp key
  
    -- When performing a distributed build you need to share a key pair
    -- (both the public and private key) with the machine you're
    -- deploying to (or from). Both machines must store the same private
    -- key at `/etc/nix/signing-key.sec` and the same public key at
    -- `/etc/nix/signing-key.pub`. The private must also be only
    -- user-readable and not group- or world-readable (i.e. `400`
    -- permissions using `chmod` notation).
    --
    -- By default, neither machine will have a key pair installed.  This script
    -- will first ensure that the remote machine has a key pair (creating one if
    -- if missing) and copy the remote key pair to the local machine.  We
    -- install the remote key pair locally on every run of this script because we
    -- do not assume that all remote machines share the same key pair.  Quite the
    -- opposite: every production machine should have a unique signing key pair.
    let privateKey = "/etc/nix/signing-key.sec"
    let publicKey  = "/etc/nix/signing-key.pub"
  
    let handler0 :: SomeException -> IO ()
        handler0 e = do
            let exceptionText = Data.Text.pack (show e)
            let msg           = [NeatInterpolation.text|
  [x] Could not ensure that the remote machine has signing keys installed
  
      Debugging tips:
  
      1. Check if you can log into the remote machine by running:
  
          $ ssh -i $key' $host
  
      2. If you can log in, then check if you have permission to `sudo` without a
         password by running the following command on the remote machine:
  
          $ sudo -n true
          $ echo $?
          0
  
      Original error: $exceptionText
  |]
            Turtle.die msg
  
    let openssl :: Turtle.Format a a
        openssl =
            "$(nix-build --no-out-link \"<nixpkgs>\" -A libressl)/bin/openssl"
    let fmt = "ssh -i "%fp%" "%s%" '"
                % "test -e "%fp%" || "
                % "sudo sh -c \""
                    % "(umask 277 && "%openssl%" genrsa -out "%fp%" 2048) && "
                    % openssl%" rsa -in "%fp%" -pubout > "%fp
                % "\""
            % "'"
    let cmd = Turtle.format fmt key host privateKey privateKey privateKey publicKey
    Control.Exception.handle handler0 (Turtle.shells cmd empty)
  
    let mirror path = Turtle.runManaged $ do
            let message = Turtle.format ("[+] Downloading: "%fp) path
            mapM_ Turtle.err (Turtle.Line.textToLines message)
  
            localPath <- Turtle.mktempfile "/tmp" "signing-key"
            let download =
                    Turtle.procs "rsync"
                        [ "--archive"
                        , "--checksum"
                        , "--rsh", Turtle.format ("ssh -i "%fp) key
                        , "--rsync-path", "sudo rsync"
                        , Turtle.format (s%":"%fp) host path
                        , Turtle.format fp localPath
                        ]
                        empty
            let handler1 :: SomeException -> IO ()
                handler1 e = do
                    let pathText      = Turtle.format fp path
                    let exceptionText = Data.Text.pack (show e)
                    let msg           = [NeatInterpolation.text|
  [x] Could not download: $pathText
  
      Debugging tips:
  
      1. Check if you can log into the remote machine by running:
  
          $ ssh -i $key' $host
  
      2. If you can log in, then check if you have permission to `sudo` without a
         password by running the following command on the remote machine:
  
          $ sudo -n true
          $ echo $?
          0
  
      3. If you can `sudo` without a password, then check if the file exists by
         running the following command on the remote machine:
  
          $ test -e $pathText
          $ echo $?
          0
  
      Original error: $exceptionText
  |]
                    Turtle.die msg
  
            liftIO (Control.Exception.handle handler1 download)
  
            new <- liftIO . Data.ByteString.Lazy.readFile . Data.Text.unpack $
              Turtle.format fp localPath
  
            old <- liftIO . Data.ByteString.Lazy.readFile . Data.Text.unpack $
              Turtle.format fp path
  
            if new == old
                then do
                    let same = Turtle.format ("[+] Unchanged: "%fp) path
                    mapM_ Turtle.err (Turtle.Line.textToLines same)
                else do
                    -- NB: path shouldn't is a FilePath and won't have any
                    -- newlines, so this should be okay
                    Turtle.err (Turtle.unsafeTextToLine $ Turtle.format ("[+] Installing: "%fp) path)
  
                    warnSudo
  
                    let install =
                            Turtle.procs "sudo"
                                [ "mv"
                                , Turtle.format fp localPath
                                , Turtle.format fp path
                                ]
                                empty
                    let handler2 :: SomeException -> IO ()
                        handler2 e = do
                            let pathText      = Turtle.format fp path
                            let exceptionText = Data.Text.pack (show e)
                            let msg           = [NeatInterpolation.text|
  [x] Could not install: $pathText
  
      Debugging tips:
  
      1. Check to see that you have permission to `sudo` by running:
  
          $ sudo true
          $ echo $?
          0
  
      Original error: $exceptionText
  |]
                            Turtle.die msg
  
                    liftIO (Control.Exception.handle handler2 install)
  
    mirror privateKey
    mirror publicKey
  
  delegateShared
      :: MonadManaged managed => OptArgs -> managed (Text, SomeException -> IO a)
  delegateShared OptArgs{..}  = do
      home <- Turtle.home
      let key'      = fromMaybe (home </> ".ssh/id_rsa") key
      let os'       = case os of [] -> [X86_64_Linux]; _ -> os
      let os''      = Data.Text.intercalate "," (Foldable.toList (fmap renderOS os'))
      let feature'  = Data.Text.intercalate "," feature
      let cores'    = fromMaybe 1 cores
  
      isDaemon <- maybe False (== "daemon") <$> Turtle.need "NIX_REMOTE"
  
      let sudo | isDaemon && canSudo cmd = "sudo"
               | otherwise               = ""
  
      host' <-
        if isDaemon && not (Data.Text.any (== '@') host)
        then do
          mUser <- Turtle.need "USER"
          case mUser of
              Nothing   -> Turtle.die [NeatInterpolation.text|
  [x] You must set the `USER` environment variable in order for `nix-delegate` to
      work in a multi-user Nix installation
  |]
              Just user -> return (user <> "@" <> host)
        else return host
  
      {-  Do a test @ssh@ command in order to prompt the user to recognize the
          host if the host is not known
  
          Use @sudo@ if we are in multi-user mode since the @root@ user will be
          initiating the build and therefore the @root@ user needs to authorize
          the known host
      -}
      Turtle.err "[+] Testing SSH access"
      if sudo == "sudo" then warnSudo else return ()
      let testSSH = s%" ssh -i "%fp%" "%s%" :"
      Turtle.shells (Turtle.format testSSH sudo key' host') Turtle.stdin
  
      liftIO (exchangeKeys key' host')
  
      let debuggingTips = [NeatInterpolation.text|
      Debugging tips:
  
      1.  Make sure that you have installed Nix:
  
          $ nix-build --version
  
      2.  Make sure that you log into a new shell after installing Nix
  |]
  
      remoteSystemsFile <- Turtle.mktempfile "/tmp" "remote-systems.conf"
      let line =
              Turtle.format
                  (s%" "%s%" "%fp%" "%d%" 1 "%s)
                  host'
                  os''
                  key'
                  cores'
                  feature'
  
      case Turtle.textToLine line of
        Just line' ->
          Turtle.output remoteSystemsFile (pure line')
        Nothing ->
          Turtle.die [NeatInterpolation.text|
  [x] The generated 'remote-systems.conf' file content contains a newline (it should not)
  
      $line
  |]
  
      loadDirectory <- Turtle.mktempdir "/tmp" "build-remote-load"
  
      mNixPath <- Turtle.need "NIX_PATH"
      nixPath  <- case mNixPath of
              Just nixPath -> return nixPath
              Nothing      -> Turtle.die [NeatInterpolation.text|
  [x] Your NIX_PATH environment variable is unset
  
      $debuggingTips
  |]
  
      let configFile = home </> ".ssh/config"
  
      configExists <- Turtle.testfile configFile
      let sshConfigFile
            | configExists = Turtle.format ("ssh-config-file="%fp%":") configFile
            | otherwise    = ""
  
      mAuthSock <- Turtle.need "SSH_AUTH_SOCK"
      let sshAuthSock = maybe "" (Turtle.format ("ssh-auth-sock="%s%":")) mAuthSock
      let nixpkgpath  = Turtle.inproc "nix-build" [ "--no-out-link", "--realise", "<nixpkgs>", "--attr", "nix" ] empty
  
      hook <- Turtle.fold nixpkgpath Foldl.head >>= \case
        Just nixpkgpath' -> do
          let nixpkgfp = Turtle.fromText $ Turtle.lineToText nixpkgpath'
          return $ Turtle.format fp (nixpkgfp </> "libexec/nix/build-remote.pl")
        Nothing ->
          Turtle.die [NeatInterpolation.text|
  [x] The 'build-remote.pl' script could not be found on your system!
  
      $debuggingTips
  |]
  
      let renderedCmd = renderCmd cmd
      let pfxcmd = Turtle.format
            (s %
             " NIX_BUILD_HOOK="%s%
             " NIX_PATH="%s%
             " NIX_REMOTE_SYSTEMS="%s%
             " NIX_CURRENT_LOAD="%s%" "%s)
            sudo
            hook
            (sshConfigFile <> sshAuthSock <> nixPath)
            (Turtle.format fp remoteSystemsFile)
            (Turtle.format fp loadDirectory)
            renderedCmd
  
      let handler2 :: SomeException -> IO a
          handler2 e = do
              let exceptionText = Data.Text.pack (show e)
              let msg           = [NeatInterpolation.text|
  [x] The subcommand you specified exited with a non-zero exit code:
  
      Original error: $exceptionText
  |]
              Turtle.die msg
  
      -- NB: path shouldn't is a FilePath and won't have any
      -- newlines, so this should be okay
      Turtle.err (Turtle.unsafeTextToLine $ Turtle.format ("[+] Running command: "%s%" "%s) sudo renderedCmd)
      Turtle.err (Turtle.unsafeTextToLine $ Turtle.format ("[+] Full command context: "%s) pfxcmd)
      return (pfxcmd, handler2)
  
  {-| Run a command with distributed builds transiently enabled
  
      This version outputs a helpful error message if the command fails
  -}
  delegate :: OptArgs -> IO ()
  delegate options = Turtle.runManaged $ do
      (command, handler) <- delegateShared options
      let build = Turtle.shells command empty
      liftIO (Control.Exception.handle handler build)
  
  {-| Run a command with distributed builds transiently enabled
  
      This version captures the output as a stream
  -}
  delegateStream :: OptArgs -> Shell Line
  delegateStream options = do
      (command, _) <- delegateShared options
      Turtle.inshell command empty
  
  warnSudo :: MonadIO io => io ()
  warnSudo = do
      exitCode <- Turtle.shell "sudo -n true 2>/dev/null" empty
  
      case exitCode of
          ExitFailure _ -> do
              Turtle.err ""
              Turtle.err "    This will prompt you for your `sudo` password"
          _             -> do
              return ()
  