path: root/fsmgr.hs

{-# OPTIONS_GHC -fno-warn-type-defaults #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE ExtendedDefaultRules  #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE NamedFieldPuns        #-}
{-# LANGUAGE NoImplicitPrelude     #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE ViewPatterns          #-}

module Main where
import           Rebase.Prelude             hiding (bool, hash, join, o, (<.>))

import           Rebase.Data.Text           (pack, unpack)
import qualified Rebase.Data.Text           as Text

import           Development.Shake          hiding (getEnv)
import           Development.Shake.Command  ()
import           Development.Shake.FilePath

import           Options.Applicative        hiding (action, command)
import qualified Options.Applicative        as Opt

import           ConfigFile
import           String
import           System.Directory           as IO
import           System.Directory           (createDirectoryIfMissing,
                                             getCurrentDirectory)
import           System.Posix.Process       (getProcessID)
import           System.Posix.Types         (CUid (..))
import           System.Posix.User          (getEffectiveUserID)

noParent :: BaseImageSpecification -> Bool
noParent (EmptyImageOfBytes _)     = True
noParent (ParentImageConfigFile _) = False
noParent (SeededImage _ _) = False

dynamicNames :: FilePath -> FilePath
dynamicNames = replace "$(kver)" (snd uname) . replace "$(karch)" (fst uname) . replace "$(debarch)" debarch

chomp :: String -> String
chomp = takeWhile (/= '\n')

debarch :: String
debarch = unsafePerformIO $ do
  Stdout out <- cmd "dpkg-architecture -q DEB_BUILD_ARCH"
  return $ chomp out

getDebianCodename :: String
getDebianCodename = unsafePerformIO $ do
  (Stdout codename) <- cmd "sh -c" [". /etc/os-release && printf '%s' \"$VERSION_CODENAME\""]
  return $ if codename == "stretch" then "buster" else codename

uname :: (String, String)
uname = unsafePerformIO $ do
  Stdout out <- cmd "sh -c" ["f=$([ -L /vmlinuz ] && readlink /vmlinuz); f=${f##*/vmlinuz-}; echo ${f:-$(uname -r)}"]
  pure $ (reverse *** reverse . tail) . break (== '-') . reverse . head . lines $ out

data AptListCfg =
  AptListCfg
    { releaseCodename :: String
    , architecture :: String
    , translationLang :: String
    }

{-
TODO: language should derive from $LC_MESSAGES as implemented by apt:

   Manual page apt.conf(5)

       Languages

           The Languages subsection controls which Translation files
           are downloaded and in which order APT tries to display
           the description-translations. APT will try to display
           the first available description in the language which is
           listed first. Languages can be defined with their short or
           long language codes.  Note that not all archives provide
           Translation files for every language - the long language
           codes are especially rare.

           The default list includes "environment" and "en".
           "environment" has a special meaning here: it will be
           replaced at runtime with the language codes extracted from
           the LC_MESSAGES environment variable. It will also ensure
           that these codes are not included twice in the list. If
           LC_MESSAGES is set to "C" only the Translation-en file (if
           available) will be used. To force APT to use no Translation
           file use the setting Acquire::Languages=none. "none" is
           another special meaning code which will stop the search for
           a suitable Translation file. This tells APT to download
           these translations too, without actually using them unless
           the environment specifies the languages. So the following
           example configuration will result in the order "en, de" in an
           English locale or "de, en" in a German one. Note that "fr"
           is downloaded, but not used unless APT is used in a French
           locale (where the order would be "fr, de, en").

               Acquire::Languages { "environment"; "de"; "en"; "none";
               "fr"; };

           Note: To prevent problems resulting from APT being executed
           in different environments (e.g. by different users or by
           other programs) all Translation files which are found in
           /var/lib/apt/lists/ will be added to the end of the list
           (after an implicit "none").

-}
language :: String
language = "en"

aptListCfg :: AptListCfg
aptListCfg = AptListCfg getDebianCodename debarch language

aptListFiles :: AptListCfg -> [FilePath]
aptListFiles AptListCfg{..} =
  ("/var/lib/apt/lists" </>) .
  (replace "Translation-en" $ "Translation-" ++ translationLang) .
  (replace "amd64" architecture) .
  (replace "buster" releaseCodename) <$> observedCorrectListForBuster
  where
    observedCorrectListForBuster =
      [ "httpredir.debian.org_debian_dists_buster_InRelease"
      , "httpredir.debian.org_debian_dists_buster_main_binary-amd64_Packages"
      , "httpredir.debian.org_debian_dists_buster_main_i18n_Translation-en"
      , "security.debian.org_dists_buster_updates_InRelease"
      , "security.debian.org_dists_buster_updates_main_binary-amd64_Packages"
      , "security.debian.org_dists_buster_updates_main_i18n_Translation-en"
      ]

buildRoot :: DiskImageConfig -> FilePath -> Action ()
buildRoot config@DiskImageConfig{..} finalOut = do
  let out        = finalOut <.> "tmp"
      mountpoint = finalOut <.> "mnt"
  cmd_ "sh -c" ["if mountpoint -q \"$0\"; then umount \"$0\"; fi", mountpoint]

  let (abortion :: IO ()) = ignoreErrors' $ do
        cmd_ "umount" [mountpoint]
        cmd_ "rmdir" [mountpoint]
        cmd_ "rm -f" [out]

  handle' abortion $ buildInitialImage config mountpoint out

  handle' abortion $ do
    {- 1. debconf -}
    forM_ debconfConfig $
      readFile' >=> liftIO . appendFile (mountpoint </> "var/cache/debconf/config.dat")
    {- 2. dpkg installs -}
    let (fmap dynamicNames -> debs, packageNames) = partitionPackages $ unpack . coerce <$> toList packages
    when (noParent initialImage || not (null packageNames)) $
      -- When there is no parent, selfstrap should install packages marked
      -- "Required" even if no packages are specified. When there is a parent,
      -- assume that this has already happened.
      cmd_ "selfstrap" (("--unpack" `consWhen` unpackOnly) ["-t", mountpoint]) (dynamicNames <$> packageNames)
    when (not $ null debs) $ do
      need debs
      cmd_ (AddEnv "DEBIAN_FRONTEND" "noninteractive")
           ["dpkg"] [if unpackOnly then "--unpack" else "--install"] ["--root", mountpoint] debs
    {- 2.5. install apt package cache -}
    when installAptLists $ do
      cmd_ "rsync -Ra" (("/./" ++) <$> aptListFiles aptListCfg) (mountpoint ++ "/")
    {- 3. binaries -}
    let go b = do
          p <- fromMaybe (fail $ "not found in ${PATH}: " ++ b) <$> liftIO (pathLocate b)
          need [p]
          cmd_ "cp -L" [p] [mountpoint </> "usr/local/bin" </> takeFileName b]
      in do
        forM_ (unpack <$> binaries) go
        forM_ (unpack <$> optionalBinaries) $ ignoreErrors . go
    {- 3.5 skel -}
    let go f = do
          homeDir <- getHomeDir
          target <- absPath mountpoint <&> (</> "etc/skel")
          cmd_ "mkdir -p" [target]
          cmd_ (Cwd homeDir) "cp -r --preserve=mode,timestamps -L --parents -t" [target] [f]
      in do
        forM_ (unpack <$> skelFiles) go
        forM_ (unpack <$> optionalSkelFiles) $ ignoreErrors . go
    {- 3.6 systemd unit files -}
    let go s = do
          target <- absPath mountpoint <&> (</> "etc/systemd/system")
          cmd_ "mkdir -p" [target]
          cmd_ "install --preserve-timestamps -m644 -t" [target] [s]
        go :: String -> Action ()
      in forM_ (unpack <$> unitFiles) go
    {- 4. custom setup commands -}
    forM_ chrootCommands $ \c -> do
      cmd_ "chroot" [mountpoint] "/bin/sh -c" [unpack c]
    cmd_ "chroot" [mountpoint] "/bin/sh -c" ["[ $(grep -c '^Package: base-files' /var/lib/dpkg/status) = 1 ]"]
    Stdout (pkgList :: String) <- cmd "chroot" [mountpoint] "dpkg -l"
    writeFileChanged (finalOut -<.> "pkgs.txt") pkgList
    {- 5. create a backup snapshot -}
    cmd_ "umount" [mountpoint]

    setupLoopDevices out
    cmd_ "mount -t btrfs -o subvol=/" [out, mountpoint]
    ignoreErrors $ cmd_ (EchoStderr False) "btrfs subvolume delete" [mountpoint </> "root~orig"]
    cmd_ "btrfs subvolume snapshot -r" [mountpoint </> "root", mountpoint </> "root~orig"]

    -- cleanup
    cmd_ "umount" [mountpoint]
    cmd_ "rmdir" [mountpoint]
    cmd_ "mv" [out, finalOut]

setupLoopDevices :: FilePath -> Action ()
setupLoopDevices out = do
  (Stdout devices) <- cmd "sh -c" ["if [ -e \"$0\" ]; then cat \"$0\"; fi ", out -<.> "devices.txt"]
  cmd_ "losetup -D"
  forM_ (lines devices) $ \d -> cmd_ "losetup -f" [d]

absPath :: MonadIO m => FilePath -> m FilePath
absPath f@('/':_) = return f
absPath f = do
  wd <- liftIO getCurrentDirectory
  return $ wd </> f

getHomeDir :: Action FilePath
getHomeDir = do
  Stdout homeDir <- cmd "sh -c" ["if [ \"$SUDO_USER\" ]; then getent passwd \"$SUDO_USER\" | cut -d: -f6; else printf \"%s\n\" \"$HOME\"; fi"]
  return homeDir

readFileOptional :: FilePath -> IO (String)
readFileOptional f = IO.doesFileExist f >>= \case
  True -> readFile f
  False -> return ""

copyParentDevices :: FilePath -> FilePath -> Action ()
copyParentDevices parent out = do
  devices <- either (const []) lines . readEither <$> liftIO (readFileOptional (parent <.> "devices.txt"))
  writeFile' (out -<.> "devices.txt") (unlines $ parent:devices)

buildInitialImage :: DiskImageConfig -> FilePath -> FilePath -> Action ()
buildInitialImage DiskImageConfig{..} mountpoint out = do
  case initialImage of
    ParentImageConfigFile f -> do
      let cwdParent = f -<.> "btrfs"
          buildDirParent = "_build" </> cwdParent
      parent <- liftIO (IO.doesFileExist cwdParent) >>= \case
        True -> return cwdParent
        False -> do
          need [buildDirParent]
          return buildDirParent
      cmd_ "cp --reflink" [parent, out]
      cmd_ "mkdir -p" [mountpoint]
      cmd_ "sh -c" ["if [ -e \"$0\" ]; then cp \"$0\" \"$1\"; fi", parent <.> "devices.txt", out -<.> "devices.txt"]

      ignoreErrors $ do cmd_ "btrfs dev scan -u"
      cmd_ "mount -t btrfs" [out] mountpoint

    SeededImage n f -> do
      let parent = f -<.> "seed.btrfs"
      need [parent]

      copyParentDevices parent out

      -- allocate new image file
      cmd_ "rm -f" [out]
      cmd_ "truncate -s" [show n, out]
      cmd_ "fallocate -l" [show n, out]

      setupLoopDevices parent
      idempotentMountImage parent mountpoint

      addImageToBtrfs out mountpoint

    EmptyImageOfBytes n -> do
      cmd_ "truncate -s" [show n] [out]
      cmd_ "mkfs.btrfs" [out]
      cmd_ "mkdir -p" [mountpoint]

      setupLoopDevices out
      cmd_ "mount -t btrfs" [out] mountpoint

      -- create new default subvolume, and then remount with it
      createDefaultSubvolume mountpoint
      cmd_ "umount" [mountpoint]
      cmd_ "mount -t btrfs" [out, mountpoint]

      cmd_ (Cwd mountpoint) "mkdir -p var/cache/apt"
      cmd_ (Cwd mountpoint) "mkdir -p var/cache/debconf"
      cmd_ (Cwd mountpoint) "btrfs subvolume create var/cache/apt/archives"

idempotentSetupLoopDev :: FilePath -> Action (Maybe String)
idempotentSetupLoopDev imageFile = do
  deleteLoopDev imageFile
  cmd_ "losetup -f" imageFile
  getLoopDev imageFile
  where
    deleteLoopDev = getLoopDev >=> mapM_ (cmd_ "losetup -d")

    getLoopDev :: FilePath -> Action (Maybe String)
    getLoopDev x = do
      Stdout r <- cmd "losetup -n -O name -j" x
      return $ guard (r /= "") >> Just r

idempotentMountImage :: FilePath -> FilePath -> Action ()
idempotentMountImage imageFile mountPoint = do
  cmd_ "mkdir -p" [mountPoint]
  mounted <- cmd "mountpoint -q" [mountPoint] <&> (== ExitSuccess)
  when mounted $ cmd_ "umount" [mountPoint]

  cmd_ "mount -o compress,ro -t btrfs" [imageFile, mountPoint]

addImageToBtrfs :: FilePath -> FilePath -> Action ()
addImageToBtrfs imageFile mountPoint = do
  blockDevice <- idempotentSetupLoopDev imageFile <&>
    fromMaybe (error "failed to set up loop device for " ++ imageFile)
  cmd_ "btrfs device add" blockDevice mountPoint
  cmd_ "mount -o remount,rw,compress" mountPoint

ignoreErrors' :: IO () -> IO ()
ignoreErrors' = flip catch (\(SomeException _) -> return ())

ignoreErrors :: Action () -> Action ()
ignoreErrors = flip actionCatch (\(SomeException _) -> return ())

partitionPackages :: [String] -> ([String], [String])
partitionPackages = partition (".deb" `isSuffixOf`)

consWhen :: a -> Bool -> [a] -> [a]
a `consWhen` c = if c then (a:) else id

strip :: String -> String
strip = unpack . Text.strip . pack

createDefaultSubvolume :: FilePath -> Action ()
createDefaultSubvolume mountpoint = do
  cmd_ (Cwd mountpoint) "btrfs subvolume create root"
  Stdout subvolIdLine <- cmd (Cwd mountpoint) "sh -c"
    ["btrfs subvolume show root | sed -n -e 's/^[ \t]*Subvolume ID:[ \t]*//p; s/.*is toplevel subvolume/5/p'"]
  let subvolId = strip subvolIdLine
  when (null subvolId) $ fail "could not obtain btrfs Subvolume ID"
  cmd_ "btrfs subvolume set-default" [strip subvolId, mountpoint]

stripSuffix :: Text -> String -> String
stripSuffix s t = fromMaybe t $ fmap unpack $ Text.stripSuffix s $ pack t

data Options = Options { optCommand :: Command }
data Command = Build BuildOpts | Chroot ChrootOpts
data BuildOpts = BuildOpts { optTarget :: String }
data ChrootOpts = ChrootOpts { optTarget :: String, optChrootCommand :: [String] }

buildOpts :: Parser Command
buildOpts = Build . BuildOpts <$> argument str idm

chrootOpts :: Parser Command
chrootOpts = fmap Chroot $ ChrootOpts <$> argument str idm <*> many (argument str idm)

main :: IO ()
main = customExecParser (prefs showHelpOnEmpty) (info opts desc) >>= run
  where
    opts :: Parser Options
    opts = Options <$> parseCommand
    desc = fullDesc <> progDesc detailed <> header "fsmgr - Debian rootfs image manager"
    detailed = unwords ["This program generates btrfs filesystem images using 'selfstrap'",
                        "which, much like 'debootstrap', creates new installations of Debian."]
    parseCommand :: Parser Command
    parseCommand = subparser $ Opt.command "build" (info' buildOpts buildDesc)
                             <> Opt.command "chroot" (info' chrootOpts chrootDesc)
    buildDesc  = ["build the image specified by the YAML config file"]
    chrootDesc = ["chroot into the mounted image.  ",
                  "Iff the chroot exits with a success value, filesystem changes will persist"]

    info' o d = info (helper <*> o) (progDesc $ unwords d)

run :: Options -> IO ()
run (Options (Build (BuildOpts target)))        = earlyFail >> shakeBuildOneImage target
run (Options (Chroot (ChrootOpts target args))) = earlyFail >> chrootImage target args

useCGroups :: Bool
useCGroups = True               -- TODO: make command-line option

handle' :: IO b -> Action a -> Action a
handle' = flip actionOnException

ourShakeOptions :: ShakeOptions
ourShakeOptions = shakeOptions {
  shakeFiles = "_build",
  shakeColor = True,
  shakeProgress = progressSimple,
  shakeVerbosity = Loud
}

chrootImage :: FilePath -> [String] -> IO ()
chrootImage target args =
  shake ourShakeOptions $ do
    shakeRules
    action $ do
      pid <- show <$> liftIO getProcessID
      let inp = target -<.> "btrfs"
      let tmp = inp <.> "tmp" <.> pid
      let mnt = tmp <.> "mnt"
      orderOnly [inp]
      cmd_ "cp --reflink=always" [inp, tmp]
      cmd_ "mkdir" [mnt]

      setupLoopDevices target
      cmd_ "mount -t btrfs" [tmp, mnt]
      let (umount :: IO ()) = do
            cmd_ "umount" [mnt]
            cmd_ "rm -f" [tmp]
      handle' umount $
        if useCGroups
          then liftIO $ cgroupChroot ("fsmgr" <.> takeFileName target) mnt args
          else cmd_ (WithStderr False) "chroot" (mnt : args)
      cmd_ "umount" [mnt]
      cmd_ "sync"
      cmd_ "mv" [tmp, inp]

shakeBuildOneImage :: FilePath -> IO ()
shakeBuildOneImage target =
  shake ourShakeOptions $ do
    want [target -<.> "btrfs"]
    shakeRules

head1 :: String -> String
head1 = lines >>> \case
  []  -> ""
  x:_ -> x

pathLocate :: String -> IO (Maybe FilePath)
pathLocate c | elem '/' c = pure $ pure c
pathLocate c | True = (getEnv "SUDO_USER" >>=) $ fmap (validatePath . head1 . fromStdout) <$> \case
  "" -> cmd "which" [c]
  u  -> do
    path <- fromStdout <$> cmd "su -" [u] "-c" ["printf %s \"$PATH\""]
    cmd (AddEnv "PATH" path) "which" [c]
  where
    validatePath "" = Nothing
    validatePath x  = Just x

infix 1 ~%>
(~%>) :: FilePattern -> (FilePath -> Action ()) -> Rules ()
filepat ~%> act = phonys f
  where
    f file | (filepat ?== file) = Just (act file)
    f _    = Nothing

cgroupChroot :: String -> FilePath -> [String] -> IO ()
cgroupChroot groupName mnt [] = cgroupChroot groupName mnt ["/bin/bash"]
cgroupChroot groupName mnt args = do
  let cgdir = "/sys/fs/cgroup/pids" </> groupName
  createDirectoryIfMissing False cgdir

  -- TODO: unshare hostname & set from /etc/hostname inside root
  cmd_ (Cwd mnt) (WithStderr False)
    "unshare --ipc --uts --cgroup --mount --pid --fork chroot ."
    "sh -exc" ["mount -t proc proc /proc; mount -t devpts devpts /dev/pts; if [ -e /etc/hostname ]; then hostname -F /etc/hostname; fi; exec \"$@\""]
    "sh" args

earlyFail :: IO ()
earlyFail = do
  CUid euid <- liftIO getEffectiveUserID
  when (euid /= 0) $ fail "you are not root"
  Stdout () <- cmd (Traced []) "which selfstrap"
  return ()

shakeRules :: Rules ()
shakeRules = do
  "_build/*.yaml.canon" %> \out -> do
    let yaml = dropDirectory1 $ dropExtension out
    need [yaml]
    cfg <- readCfg yaml
    writeFileChanged out (show cfg)
  "_build/*.btrfs" %> \out -> do
    let cfgFile = (out -<.> "yaml.canon")
    need [cfgFile]
    cfg <- readEither <$> readFile' cfgFile
    either (error . (("Error parsing file: " ++ cfgFile ++ ": ") ++)) (flip buildRoot out) cfg
  priority 2 $ "*.seed.btrfs" %> \out -> do
    let tmp = out <.> "tmp"
        inp = dropExtension out -<.> ".btrfs"
    need [inp]
    cmd_ "cp --reflink=always" [inp, tmp]
    setupLoopDevices ("_build" </> inp <.> "btrfs")
    cmd_ "btrfs-shrink" [tmp]
    cmd_ "btrfstune -f -S1" [tmp]
    cmd_ (WithStderr False) "mv -i" [tmp, out]

  priority 1 $ "*.btrfs" %> \out -> do
    need ["_build" </> out]
    -- WithStderr False needed for `cp` to interact with the tty
    cmd_ (WithStderr False) "cp --reflink=always -i" ["_build" </> out, out]