rofi-extras/app/rofi-dev.hs

{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE NoGeneralizedNewtypeDeriving #-}

--------------------------------------------------------------------------------
-- rofi-dev - a rofi prompt for mountable devices
--
-- Like all "mount helpers" this is basically a wrapper for low-level utilities
-- the mount things from the command line. It also creates/destroys mountpoint
-- paths given a specific location for such mountpoints.

module Main (main) where

import Bitwarden.Internal
import qualified Data.Text.IO as TI
import Dhall hiding (maybe, sequence, void)
import Dhall.TH
import RIO
import RIO.Directory
import qualified RIO.List as L
import qualified RIO.Map as M
import qualified RIO.NonEmpty as NE
import qualified RIO.Text as T
import Rofi.Command
import System.Console.GetOpt
import System.Environment
import System.FilePath.Posix
import System.Posix.User (getEffectiveUserName)
import System.Process

--------------------------------------------------------------------------------
-- Static device configuration (dhall)

makeHaskellTypesWith
  (defaultGenerateOptions {generateToDhallInstance = False})
  [ MultipleConstructors "PasswordConfig" "(./dhall/rofi-dev.dhall).PasswordConfig"
  , MultipleConstructors "DataConfig" "(./dhall/rofi-dev.dhall).DataConfig"
  , SingleConstructor "TreeMap" "TreeMap" "(./dhall/rofi-dev.dhall).TreeMap"
  , SingleConstructor "SecretMap" "SecretMap" "(./dhall/rofi-dev.dhall).SecretMap"
  , SingleConstructor "StaticConfig" "StaticConfig" "(./dhall/rofi-dev.dhall).StaticConfig"
  , SingleConstructor "PromptConfig" "PromptConfig" "(./dhall/rofi-dev.dhall).PromptConfig"
  , SingleConstructor "TreeConfig" "TreeConfig" "(./dhall/rofi-dev.dhall).TreeConfig"
  , SingleConstructor "DeviceConfig" "DeviceConfig" "(./dhall/rofi-dev.dhall).DeviceConfig"
  , SingleConstructor "SecretConfig" "SecretConfig" "(./dhall/rofi-dev.dhall).SecretConfig"
  , SingleConstructor "MountConfig" "MountConfig" "(./dhall/rofi-dev.dhall).MountConfig"
  , SingleConstructor "BitwardenConfig" "BitwardenConfig" "(./dhall/rofi-dev.dhall).BitwardenConfig"
  , SingleConstructor "VeracryptData" "VeracryptData" "(./dhall/rofi-dev.dhall).VeracryptData"
  , SingleConstructor "CIFSData" "CIFSData" "(./dhall/rofi-dev.dhall).CIFSData"
  , SingleConstructor "CIFSOpts" "CIFSOpts" "(./dhall/rofi-dev.dhall).CIFSOpts"
  , SingleConstructor "SSHFSData" "SSHFSData" "(./dhall/rofi-dev.dhall).SSHFSData"
  ]

main :: IO ()
main = getArgs >>= parse

parse :: [String] -> IO ()
parse args = case getOpt Permute options args of
  (o, n, []) -> runMounts $ L.foldl (flip id) (defaultOpts (fmap T.pack n)) o
  (_, _, errs) -> TI.putStrLn $ T.pack $ concat errs ++ usageInfo h options
  where
    h = "Usage: rofi-dev [OPTIONS] [-- ROFI-OPTIONS]"
    defaultOpts r =
      Opts
        { optsConfig = Nothing
        , optsAlias = Nothing
        , optsUnmount = False
        , optsRofiArgs = r
        }

options :: [OptDescr (Opts -> Opts)]
options =
  [ Option
      ['c']
      ["config"]
      (ReqArg (\s m -> m {optsConfig = Just s}) "CONF")
      "The path to the config file"
  , Option
      ['m']
      ["mount"]
      (ReqArg (\s m -> m {optsAlias = Just $ T.pack s}) "ALIAS")
      "Mount the device specified by ALIAS directly"
  , Option
      ['u']
      ["unmount"]
      (NoArg (\m -> m {optsUnmount = True}))
      "Unmount the device specified by ALIAS instead of mounting it."
  ]

data Opts = Opts
  { optsConfig :: Maybe FilePath
  , optsAlias :: Maybe T.Text
  , optsUnmount :: Bool
  , optsRofiArgs :: [T.Text]
  }
  deriving (Show)

--------------------------------------------------------------------------------
-- Main prompt
--
-- This command will have one Rofi prompt and will display all available
-- mounts grouped by device type (eg removable, sshfs, cifs, etc). I like
-- pretty things, so ensure the entries are aligned properly as well

runMounts :: Opts -> IO ()
runMounts opts = do
  static <- join <$> traverse parseStaticConfig (optsConfig opts)
  defaultTmpPath <- ("/tmp/media" </>) <$> getEffectiveUserName
  let tmpPath = fromMaybe defaultTmpPath $ (fmap T.unpack . scTmpPath) =<< static
  let staticDevs = maybe M.empty (M.fromList . fmap (\(TreeMap k v) -> (k, v)) . scDevices) static
  let verbose = fromMaybe False $ scVerbose =<< static
  let mountconf =
        MountConf
          { mountconfVolatilePath = tmpPath
          , mountconfRofiArgs = optsRofiArgs opts
          , mountconfStaticDevs = staticDevs
          , mountconfVerbose = verbose
          }
  let byAlias = mountByAlias $ optsUnmount opts
  let byPrompt = runPrompt =<< getGroups
  runRIO mountconf $ maybe byPrompt byAlias $ optsAlias opts

parseStaticConfig :: FilePath -> IO (Maybe StaticConfig)
parseStaticConfig p = do
  res <- tryIO $ inputFile auto p
  case res of
    Left e -> TI.putStrLn (T.pack $ show e) >> return Nothing
    Right c -> return $ Just c

runPrompt :: HasRofiConf c => [RofiGroup c] -> RIO c ()
runPrompt gs =
  selectAction $
    emptyMenu
      { groups = gs
      , prompt = Just "Select Device"
      }

getGroups :: RofiMountIO [RofiGroup MountConf]
getGroups = do
  actions <- sequence [getStaticActions, getRemovableActions, getMTPActions]
  return $
    (++ [metaActions]) $
      fmap mkGroup $
        NE.groupAllWith fst $
          concat actions
  where
    metaActions =
      titledGroup "Meta Actions" $
        toRofiActions [("  Dismount All", dismountAll)]

dismountAll :: RofiMountIO ()
dismountAll = do
  umount =<< asks (configToTree' . mountconfStaticDevs)
  umount =<< getRemovableDevices
  umount =<< getMTPDevices
  return ()
  where
    umount :: Mountable a => [a] -> RofiMountIO ()
    umount = mapM_ (`mountMaybe` True)

mountByAlias :: Bool -> T.Text -> RofiMountIO ()
mountByAlias unmountFlag alias = do
  static <- asks mountconfStaticDevs
  mapM_ (`mountMaybe` unmountFlag) $ configToTree static <$> M.lookup alias static

mkGroup :: NE.NonEmpty (Header, ProtoAction) -> RofiGroup MountConf
mkGroup as = titledGroup h $ toRofiActions $ NE.toList $ alignEntries $ snd <$> as
  where
    h = (T.pack $ show $ fst $ NE.head as)

alignSep :: T.Text
alignSep = " | "

alignEntries :: NE.NonEmpty (ProtoAction) -> NE.NonEmpty (T.Text, RofiMountIO ())
alignEntries ps = NE.zip (align es) as
  where
    (es, as) = NE.unzip $ fmap (\(ProtoAction e a) -> (e, a)) ps
    align =
      fmap (T.intercalate alignSep . NE.toList)
        . NE.transpose
        . fmap1 padAll
        . NE.transpose
    fmap1 f (x :| xs) = (f x) :| xs
    padAll xs = let m = maxNE $ fmap T.length xs in fmap (rpad m ' ') xs
    maxNE (x :| []) = x
    maxNE (x :| (y : ys)) = maxNE $ (max x y) :| ys

rpad :: Int -> Char -> T.Text -> T.Text
rpad n c s = T.append s $ T.replicate (n - T.length s) $ T.singleton c

--------------------------------------------------------------------------------
-- Global config used in the reader monad stack

data MountConf = MountConf
  { mountconfVolatilePath :: FilePath
  , mountconfRofiArgs :: [T.Text]
  , mountconfStaticDevs :: M.Map T.Text TreeConfig
  , mountconfVerbose :: Bool
  }

-- deriving (Show)

instance HasRofiConf MountConf where
  defArgs MountConf {mountconfRofiArgs = a} = a

--------------------------------------------------------------------------------
-- Mountable typeclass
--
-- Class to provide common interface for anything that can be mounted.

data MountState = Unmounted | Mounted | Partial deriving (Show, Eq)

mountedState :: MountState -> Bool
mountedState Mounted = True
mountedState _ = False

class Mountable a where
  -- | Mount the given type (or dismount if False is passed)
  mount :: a -> Bool -> RofiMountIO MountResult

  mountMaybe :: a -> Bool -> RofiMountIO ()
  mountMaybe dev mountFlag = do
    let lab = getLabel dev
    mounted <- isMounted dev
    verbose <- asks mountconfVerbose
    if
        | mountFlag == mounted -> do
            r <- mount dev mountFlag
            io $ notifyMountResult mounted lab r
        | verbose ->
            io $ notify IconInfo (T.append lab " already mounted") Nothing
        | otherwise ->
            return ()

  -- | Check if the mounting utilities are present
  allInstalled :: a -> RofiMountIO Bool

  -- | Return a string representing the label of the device
  getLabel :: a -> T.Text

  -- | Determine if the given type is mounted or not
  isMounted :: a -> RofiMountIO Bool
  isMounted dev = mountedState <$> mountState dev

  mountState :: a -> RofiMountIO MountState

--------------------------------------------------------------------------------
-- Actionable typeclass
--
-- Class to provide common interface for anything that can be presented in the
-- Rofi menu as an action. Note that this must be separate from the Mountable
-- class above because some devices are represented as trees, and displaying
-- these trees in the rofi menu only requires that the tree itself be presented
-- and not its subcomponents.

class Mountable a => Actionable a where
  -- | Return a string to go in the Rofi menu for the given type
  fmtEntry :: a -> NE.NonEmpty T.Text
  fmtEntry d = (getLabel d :| [])

  groupHeader :: a -> Header

  -- | Given a mountable type, return a rofi action (string to go in the
  -- Rofi prompt and an action to perform when it is selected)
  mkAction :: a -> RofiMountIO (Header, ProtoAction)
  mkAction dev = do
    m <- mountState dev
    i <- allInstalled dev
    let h = groupHeader dev
    let action = when i $ mountMaybe dev $ mountedState m
    let entry = case fmtEntry dev of
          (e :| es) -> (T.append (mountedPrefix m i) e) :| es
    return (h, ProtoAction entry action)
    where
      mountedPrefix _ False = "! "
      mountedPrefix Unmounted True = "  "
      mountedPrefix Mounted True = "* "
      mountedPrefix Partial True = "- "

mountableToAction
  :: Actionable a
  => RofiMountIO [a]
  -> RofiMountIO [(Header, ProtoAction)]
mountableToAction ms = mapM mkAction =<< ms

type RofiMountIO a = RIO MountConf a

-- headers appear in the order listed here (per Enum)
data Header
  = CIFSHeader
  | SSHFSHeader
  | VeracryptHeader
  | RemovableHeader
  | MTPFSHeader
  deriving (Ord, Enum, Eq)

instance Show Header where
  show h = case h of
    CIFSHeader -> suffix "CIFS"
    SSHFSHeader -> suffix "SSHFS"
    VeracryptHeader -> suffix "Veracrypt"
    RemovableHeader -> suffix "Removable"
    MTPFSHeader -> suffix "MTPFS"
    where
      suffix = (++ " Devices")

data ProtoAction = ProtoAction (NE.NonEmpty T.Text) (RofiMountIO ())

--------------------------------------------------------------------------------
-- Static devices trees

-- Static devices as defined in the config file may declare dependencies on
-- other static devices, and thus are best represented as a tree. Note that the
-- tree is both Actionable and Mountable, where each node in the tree is only
-- Mountable; this is because trees need to be displayed and chosen in the Rofi
-- menu.

data Tree a = Tree a [Tree a] deriving (Eq, Show)

type StaticConfigTree = Tree DeviceConfig

instance Mountable a => Mountable (Tree a) where
  mount (Tree p cs) False = mapM_ (`mountMaybe` False) cs >> mount p False
  mount (Tree p _) True = mount p True

  mountState (Tree p _) = mountState p

  allInstalled (Tree p cs) = do
    res <- and <$> mapM allInstalled cs
    if res then allInstalled p else return res

  getLabel (Tree p _) = getLabel p

instance Actionable (Tree DeviceConfig) where
  fmtEntry (Tree p@DeviceConfig {deviceData = d} _) = (getLabel p :| [target d])
    where
      target (CIFSConfig (CIFSData {cifsRemote = r})) = r
      target (SSHFSConfig (SSHFSData {sshfsRemote = r})) = r
      target (VeracryptConfig (VeracryptData {vcVolume = v})) = v

  groupHeader (Tree DeviceConfig {deviceData = d} _) =
    case d of
      CIFSConfig {} -> CIFSHeader
      SSHFSConfig {} -> SSHFSHeader
      VeracryptConfig {} -> VeracryptHeader

configToTree' :: M.Map T.Text TreeConfig -> [StaticConfigTree]
configToTree' devMap = configToTree devMap <$> M.elems devMap

configToTree :: M.Map T.Text TreeConfig -> TreeConfig -> StaticConfigTree
configToTree devMap TreeConfig {tcParent = p, tcChildren = c} =
  -- TODO wut?
  Tree p $ fmap go id c
  where
    go ds = configToTree devMap <$> mapMaybe (`M.lookup` devMap) ds

--------------------------------------------------------------------------------
-- Static devices
--
-- This is complex because there may be multiple classes of static devices
-- in the config file, and each device may depend on another device that is
-- a different class (eg sshfs on cifs). I deal with this by abstracting the
-- differences between each class in a sum-record type; in this way the
-- processing "splits" and "converges" entirely in this function, so nothing
-- outside of these needs to be aware of these different classes.

instance Mountable DeviceConfig where
  mount DeviceConfig {deviceMount = m, deviceData = devData} False = do
    m' <- getAbsMountpoint m
    withTmpMountDir m' $
      io $
        case devData of
          SSHFSConfig (SSHFSData {sshfsRemote = r, sshfsPassword = p}) ->
            mountSSHFS m' p r
          CIFSConfig
            ( CIFSData
                { cifsRemote = r
                , cifsSudo = s
                , cifsPassword = p
                , cifsOpts = o
                }
              ) ->
              mountCIFS s r m' o p
          VeracryptConfig
            ( VeracryptData
                { vcPassword = p
                , vcVolume = v
                }
              ) ->
              mountVeracrypt m' p v
  mount DeviceConfig {deviceMount = m, deviceData = d} True = do
    m' <- getAbsMountpoint m
    runAndRemoveDir m' $ io $ case d of
      CIFSConfig (CIFSData {cifsSudo = s}) -> runMountSudoMaybe s "umount" [T.pack m']
      VeracryptConfig _ -> runVeraCrypt ["-d", T.pack m'] ""
      _ -> runMount "umount" [T.pack m'] ""

  allInstalled DeviceConfig {deviceData = devData} =
    io $
      isJust
        <$> findExecutable (exe devData)
    where
      exe SSHFSConfig {} = "sshfs"
      exe CIFSConfig {} = "mount.cifs"
      exe VeracryptConfig {} = "veracrypt"

  mountState DeviceConfig {deviceMount = m, deviceData = d} = do
    -- mountState DeviceConfig{ deviceMount = m } = do
    case d of
      VeracryptConfig {} -> veracryptMountState m
      _ -> do
        b <- (io . isDirMounted) =<< getAbsMountpoint m
        return $ if b then Mounted else Unmounted

  getLabel
    DeviceConfig
      { deviceMount = MountConfig {mpPath = p, mpLabel = l}
      } = fromMaybe (T.pack $ takeFileName $ T.unpack p) l

mountSSHFS :: FilePath -> Maybe PasswordConfig -> T.Text -> IO MountResult
mountSSHFS mountpoint pwdConfig remote =
  withPasswordGetter pwdConfig (run ["-o", "password_stdin"]) $ run [] ""
  where
    run other = runMount "sshfs" (other ++ [remote, T.pack mountpoint])

mountCIFS
  :: Bool
  -> T.Text
  -> FilePath
  -> Maybe CIFSOpts
  -> Maybe PasswordConfig
  -> IO MountResult
mountCIFS useSudo remote mountpoint opts pwdConfig =
  withPasswordGetter pwdConfig runPwd run
  where
    run = runMountSudoMaybe useSudo "mount.cifs" args
    runPwd p = runMountSudoMaybe' useSudo "mount.cifs" args [("PASSWD", p)]
    args = [remote, T.pack mountpoint] ++ maybe [] (\o -> ["-o", fromCIFSOpts o]) opts

fromCIFSOpts :: CIFSOpts -> T.Text
fromCIFSOpts o = T.intercalate "," $ mapMaybe concatMaybe fs
  where
    fs =
      [ ("username", cifsoptsUsername)
      , ("workgroup", cifsoptsWorkgroup)
      , ("uid", fmap (T.pack . show) . cifsoptsUID)
      , ("gid", fmap (T.pack . show) . cifsoptsGID)
      , ("iocharset", cifsoptsIocharset)
      ]
    concatMaybe (k, f) = (\v -> T.concat [k, "=", v]) <$> f o

mountVeracrypt :: FilePath -> Maybe PasswordConfig -> T.Text -> IO MountResult
mountVeracrypt mountpoint pwdConfig volume =
  withPasswordGetter pwdConfig (runVeraCrypt (args ++ ["--stdin"])) $
    runVeraCrypt args ""
  where
    args = [volume, T.pack mountpoint]

-- NOTE: the user is assumed to have added themselves to the sudoers file so
-- that this command will work
runVeraCrypt :: [T.Text] -> T.Text -> IO MountResult
runVeraCrypt args = runMount "sudo" (defaultArgs ++ args)
  where
    defaultArgs = ["/usr/bin/veracrypt", "--text", "--non-interactive"]

veracryptMountState :: MountConfig -> RofiMountIO MountState
veracryptMountState mc = do
  mp <- getAbsMountpoint mc
  primary <- io $ lookupSpec mp
  aux <- io $ fmap join $ mapM lookupSpec $ auxPath =<< primary
  return $ case (primary, aux) of
    (Just _, Just _) -> Mounted
    (Nothing, Nothing) -> Unmounted
    _ -> Partial
  where
    -- TODO don't hardcode the tmp directory
    auxPath = fmap (\i -> "/tmp/.veracrypt_aux_mnt" ++ [i]) . vcIndex
    vcIndex spec = case T.uncons $ T.reverse spec of
      -- TODO what if I have more than one digit?
      Just (i, _) -> if i `elem` ['0' .. '9'] then Just i else Nothing
      _ -> Nothing

getAbsMountpoint :: MountConfig -> RofiMountIO FilePath
getAbsMountpoint MountConfig {mpPath = m} =
  asks $ flip appendRoot (T.unpack m) . mountconfVolatilePath

getStaticActions :: RofiMountIO [(Header, ProtoAction)]
getStaticActions = mountableToAction $ asks $ configToTree' . mountconfStaticDevs

--------------------------------------------------------------------------------
-- Password-getting functions for static devices

type PasswordGetter = IO (Maybe T.Text)

runSecret :: M.Map T.Text T.Text -> PasswordGetter
runSecret kvs = readCmdSuccess "secret-tool" ("lookup" : kvs') ""
  where
    kvs' = concatMap (\(k, v) -> [k, v]) $ M.toList kvs

runBitwarden :: T.Text -> PasswordGetter
runBitwarden pname =
  ((password . login) <=< L.find (\i -> name i == pname))
    <$> getItems

runPromptLoop :: Natural -> PasswordGetter -> PasswordGetter
runPromptLoop n pwd = do
  res <- pwd
  if isNothing res
    then if n <= 0 then return Nothing else runPromptLoop (n - 1) pwd
    else return res

-- configToPwd :: PasswordConfig -> PasswordGetter
-- configToPwd PasswordConfig
--   { passwordBitwarden = b
--   , passwordLibSecret = s
--   , passwordPrompt = p
--   } =
--     getBW b `runMaybe` getLS s `runMaybe` getPrompt p
--   where
--     getBW (Just BitwardenConfig{ bwKey = k, bwTries = n }) =
--       runPromptLoop n $ runBitwarden k
--     getBW _ = return Nothing
--     getLS = maybe (return Nothing) (runSecret . secretAttributes)
--     getPrompt = maybe (return Nothing) (flip runPromptLoop readPassword . promptTries)
--     runMaybe x y = (\r -> if isNothing r then y else return r) =<< x

configToPwd :: PasswordConfig -> PasswordGetter
configToPwd (PwdBW (BitwardenConfig {bwKey = k, bwTries = n})) =
  runPromptLoop n $ runBitwarden k
configToPwd (PwdLS s) = runSecret $ M.fromList $ fmap (\(SecretMap k v) -> (k, v)) $ secretAttributes s
configToPwd (PwdPr p) = flip runPromptLoop readPassword $ promptTries p

withPasswordGetter
  :: Maybe PasswordConfig
  -> (T.Text -> IO MountResult)
  -> IO MountResult
  -> IO MountResult
withPasswordGetter (Just pwdConfig) runPwd _ =
  maybe (return $ MountError "Password could not be obtained") runPwd
    =<< configToPwd pwdConfig
withPasswordGetter Nothing _ run = run

--------------------------------------------------------------------------------
-- Removable devices
--
-- A device which can be removed (such as a flash drive). These are distinct
-- from any device in the static configuration in that they only have device
-- addresses (eg in /dev) and labels.

data Removable = Removable
  { removablePath :: T.Text
  , removableLabel :: T.Text
  }
  deriving (Eq, Show)

instance Mountable Removable where
  mount Removable {removablePath = d} m =
    io $ runMount "udisksctl" [c, "-b", d] ""
    where
      c = if m then "unmount" else "mount"

  allInstalled _ = fmap isJust $ io $ findExecutable "udisksctl"

  mountState Removable {removablePath = d} = do
    s <- elem d <$> io curDeviceSpecs
    return $ if s then Mounted else Unmounted

  getLabel Removable {removableLabel = l} = l

instance Actionable Removable where
  fmtEntry Removable {removablePath = d, removableLabel = l} = (l :| [d])

  groupHeader _ = RemovableHeader

-- | Return list of possible rofi actions for removable devices
-- A 'removable device' is defined as a hotplugged device with a filesystem as
-- reported by 'lsblk'. If the LABEL does not exist on the filesystem, the
-- label shown on the prompt will be 'SIZE Volume' where size is the size of
-- the device
getRemovableDevices :: RIO c [Removable]
getRemovableDevices =
  fromLines toDev . T.lines . T.pack
    <$> io (readProcess "lsblk" ["-n", "-r", "-o", columns] "")
  where
    columns = "FSTYPE,HOTPLUG,PATH,LABEL,SIZE"
    -- can't use 'words' here since it will drop spaces in the front
    toDev line = case T.split (== ' ') line of
      ("" : _) -> Nothing
      [_, "1", d, "", s] -> mk d $ T.append s " Volume"
      [_, "1", d, l, _] -> mk d l
      _ -> Nothing
    mk d l = Just $ Removable {removablePath = d, removableLabel = l}

getRemovableActions :: RofiMountIO [(Header, ProtoAction)]
getRemovableActions = mountableToAction getRemovableDevices

--------------------------------------------------------------------------------
-- MTP devices

mtpExe :: FilePath
mtpExe = "jmtpfs"

data MTPFS = MTPFS
  { mtpfsBus :: T.Text
  , mtpfsDevice :: T.Text
  , mtpfsMountpoint :: FilePath
  , mtpfsDescription :: T.Text
  }
  deriving (Eq, Show)

instance Mountable MTPFS where
  mount MTPFS {mtpfsBus = b, mtpfsDevice = n, mtpfsMountpoint = m} False = do
    -- TODO add autodismount to options
    let dev = T.concat ["-device=", b, ",", n]
    withTmpMountDir m $ io $ runMount (T.pack mtpExe) [dev, T.pack m] ""
  mount MTPFS {mtpfsMountpoint = m} True =
    runAndRemoveDir m $ io $ runMount "umount" [T.pack m] ""

  -- \| return True always since the list won't even show without jmtpfs
  allInstalled _ = return True

  mountState MTPFS {mtpfsMountpoint = m} = do
    s <- io $ isDirMounted m
    return $ if s then Mounted else Unmounted

  getLabel = mtpfsDescription

-- | Return list of all available MTP devices
getMTPDevices :: RofiMountIO [MTPFS]
getMTPDevices = do
  i <- io mtpExeInstalled
  if i then go else return []
  where
    go = do
      dir <- asks mountconfVolatilePath
      res <- io $ readProcess mtpExe ["-l"] ""
      return $ fromLines (toDev dir) $ toDevList $ T.pack res
    toDevList =
      reverse
        . L.takeWhile (not . T.isPrefixOf "Available devices")
        . L.reverse
        . T.lines
    toDev dir s = case L.filter (== " ") $ T.split (== ',') s of
      [busNum, devNum, _, _, desc, vendor] ->
        let d = T.unwords [vendor, desc]
         in Just $
              MTPFS
                { mtpfsBus = busNum
                , mtpfsDevice = devNum
                , mtpfsMountpoint = dir </> canonicalize (T.unpack d)
                , mtpfsDescription = d
                }
      _ -> Nothing
    canonicalize = mapMaybe repl
    repl c
      | c `elem` ("\"*/:<>?\\|" :: [Char]) = Nothing
      | c == ' ' = Just '-'
      | otherwise = Just c

getMTPActions :: RofiMountIO [(Header, ProtoAction)]
getMTPActions = mountableToAction getMTPDevices

mtpExeInstalled :: IO Bool
mtpExeInstalled = isJust <$> findExecutable mtpExe

instance Actionable MTPFS where
  fmtEntry d = (getLabel d :| [])

  groupHeader _ = MTPFSHeader

--------------------------------------------------------------------------------
-- Notifications

data NotifyIcon = IconError | IconInfo

instance Show NotifyIcon where
  show IconError = "dialog-error-symbolic"
  show IconInfo = "dialog-information-symbolic"

notifyMountResult :: Bool -> T.Text -> MountResult -> IO ()
notifyMountResult mounted label result = case result of
  MountError e -> notify IconError (T.unwords ["Failed", "to", verb, label]) $ Just e
  MountSuccess -> notify IconInfo (T.concat ["Successfully ", verb, "ed ", label]) Nothing
  where
    verb = if mounted then "unmount" else "mount" :: T.Text

notify :: NotifyIcon -> T.Text -> Maybe T.Text -> IO ()
notify icon summary body =
  void $
    spawnProcess "notify-send" $
      maybe args (\b -> args ++ [b]) $
        fmap T.unpack body
  where
    args = ["-i", show icon, T.unpack summary]

--------------------------------------------------------------------------------
-- Mount commands

data MountResult = MountSuccess | MountError T.Text deriving (Show, Eq)

runMount :: T.Text -> [T.Text] -> T.Text -> IO MountResult
runMount cmd args stdin_ = eitherToMountResult <$> readCmdEither cmd args stdin_

runMount' :: T.Text -> [T.Text] -> T.Text -> [(T.Text, T.Text)] -> IO MountResult
runMount' cmd args stdin_ environ =
  eitherToMountResult
    <$> readCmdEither' cmd args stdin_ environ

runMountSudoMaybe :: Bool -> T.Text -> [T.Text] -> IO MountResult
runMountSudoMaybe useSudo cmd args =
  runMountSudoMaybe' useSudo cmd args []

runMountSudoMaybe' :: Bool -> T.Text -> [T.Text] -> [(T.Text, T.Text)] -> IO MountResult
runMountSudoMaybe' useSudo cmd args environ =
  maybe
    (runMount' cmd args "" environ)
    (\r -> runSudoMount' r cmd args environ)
    =<< if useSudo then readPassword' "Sudo Password" else return Nothing

-- TODO untested
-- runSudoMount :: T.Text -> T.Text -> [T.Text] -> T.Text -> IO MountResult
-- runSudoMount rootpass cmd args stdin = runSudoMount' rootpass cmd args stdin []

runSudoMount' :: T.Text -> T.Text -> [T.Text] -> [(T.Text, T.Text)] -> IO MountResult
runSudoMount' rootpass cmd args environ = runMount "sudo" args' rootpass
  where
    args' = ["-S"] ++ environ' ++ [cmd] ++ args
    environ' = fmap (\(k, v) -> T.concat [k, "=", v]) environ

eitherToMountResult :: Either (Int, T.Text, T.Text) T.Text -> MountResult
eitherToMountResult (Right _) = MountSuccess
eitherToMountResult (Left (_, _, e)) = MountError e

--------------------------------------------------------------------------------
-- Low-level mount functions

mountMap :: IO (M.Map FilePath T.Text)
mountMap = do
  parseFile <$> readFileUtf8 "/proc/mounts"
  where
    parseFile = M.fromList . mapMaybe (parseLine . T.words) . T.lines
    -- none of these should fail since this file format will never change
    parseLine [spec, mountpoint, _, _, _, _] = Just (T.unpack mountpoint, spec)
    parseLine _ = Nothing

curDeviceSpecs :: IO [T.Text]
curDeviceSpecs = M.elems <$> mountMap

curMountpoints :: IO [FilePath]
curMountpoints = M.keys <$> mountMap

lookupSpec :: FilePath -> IO (Maybe T.Text)
lookupSpec mountpoint = M.lookup mountpoint <$> mountMap

-- ASSUME the base mount path will always be created because
-- 'createDirectoryIfMissing' will make parents if missing, and that removing
-- all the directories will leave the base mount path intact regardless of if it
-- was present before doing anything (which matters here since I'm putting the
-- base path in /tmp, so all this is saying is that umounting everything will
-- leave /tmp/media/USER without removing all the way down to /tmp)

rmDirOnMountError :: FilePath -> RofiMountIO MountResult -> RofiMountIO MountResult
rmDirOnMountError d f = do
  res <- f
  unless (res == MountSuccess) $ rmDirMaybe d
  return res

-- | Run a mount command and create the mountpoint if it does not exist, and
-- remove the mountpoint if a mount error occurs
withTmpMountDir :: FilePath -> RofiMountIO MountResult -> RofiMountIO MountResult
withTmpMountDir m =
  rmDirOnMountError m
    . bracketOnError_ (mkDirMaybe m) (rmDirMaybe m)

-- | Run an unmount command and remove the mountpoint if no errors occur
runAndRemoveDir :: FilePath -> RofiMountIO MountResult -> RofiMountIO MountResult
runAndRemoveDir m f = do
  res <- catch f (return . MountError . (T.pack . displayException :: SomeException -> T.Text))
  when (res == MountSuccess) $ rmDirMaybe m
  return res

mkDirMaybe :: FilePath -> RofiMountIO ()
mkDirMaybe fp = whenInMountDir fp $ io $ createDirectoryIfMissing True fp

rmDirMaybe :: FilePath -> RofiMountIO ()
rmDirMaybe fp =
  whenInMountDir fp $
    unlessMountpoint fp $
      asks mountconfVolatilePath >>= io . rmUntil fp
  where
    rmUntil cur target = unless (target == cur) $ do
      removePathForcibly cur
      rmUntil (takeDirectory cur) target

whenInMountDir :: FilePath -> RofiMountIO () -> RofiMountIO ()
whenInMountDir fp f = do
  mDir <- asks mountconfVolatilePath
  when (mDir `L.isPrefixOf` fp) f

unlessMountpoint :: MonadIO m => FilePath -> m () -> m ()
unlessMountpoint fp f = do
  mounted <- io $ isDirMounted fp
  unless mounted f

isDirMounted :: FilePath -> IO Bool
isDirMounted fp = elem fp <$> curMountpoints

--------------------------------------------------------------------------------
-- Other functions

fromLines :: (T.Text -> Maybe a) -> [T.Text] -> [a]
fromLines f = mapMaybe (f . T.strip)

-- TODO this exists somewhere...
-- splitBy :: Char -> T.Text -> [T.Text]
-- splitBy delimiter = T.foldr f [[]]
--   where
--     f _ [] = []
--     f c l@(x : xs)
--       | c == delimiter = [] : l
--       | otherwise = (c : x) : xs

appendRoot :: FilePath -> FilePath -> FilePath
appendRoot root path = if isRelative path then root </> path else path