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 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 Rofi.IO
import System.Console.GetOpt
import System.FilePath.Posix
import System.Posix.User (getEffectiveUserName)
import System.Process
import UnliftIO.Environment

--------------------------------------------------------------------------------
-- 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 = runSimpleApp $ getArgs >>= parse

parse :: [String] -> RIO SimpleApp ()
parse args = case getOpt Permute options args of
  (o, n, []) -> runMounts $ L.foldl (flip id) (defaultOpts (fmap T.pack n)) o
  (_, _, errs) ->
    logError $
      displayBytesUtf8 $
        encodeUtf8 $
          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 -> RIO SimpleApp ()
runMounts opts = do
  static <- join <$> traverse parseStaticConfig (optsConfig opts)
  defaultTmpPath <- ("/tmp/media" </>) <$> liftIO 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 e =
        MountConf
          { mountconfVolatilePath = tmpPath
          , mountconfRofiArgs = optsRofiArgs opts
          , mountconfStaticDevs = staticDevs
          , mountconfVerbose = verbose
          , mountconfEnv = e
          }
  let byAlias = mountByAlias $ optsUnmount opts
  let byPrompt = runPrompt =<< getGroups
  mapRIO mountconf $ maybe byPrompt byAlias $ optsAlias opts

parseStaticConfig
  :: (HasLogFunc c, MonadReader c m, MonadUnliftIO m)
  => FilePath
  -> m (Maybe StaticConfig)
parseStaticConfig p = do
  res <- tryIO $ liftIO $ inputFile auto p
  case res of
    Left e -> do
      logError $ displayBytesUtf8 $ encodeUtf8 $ 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 :: MIO [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 :: MIO ()
dismountAll = do
  umount =<< asks (configToTree' . mountconfStaticDevs)
  umount =<< getRemovableDevices
  umount =<< getMTPDevices
  return ()
  where
    umount :: Mountable a => [a] -> MIO ()
    umount = mapM_ (`mountMaybe` True)

mountByAlias :: Bool -> T.Text -> MIO ()
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, MIO ())
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
  , mountconfEnv :: !SimpleApp
  }

-- deriving (Show)

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

instance HasLogFunc MountConf where
  logFuncL = lens mountconfEnv (\x y -> x {mountconfEnv = y}) . logFuncL

--------------------------------------------------------------------------------
-- 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 -> MIO MountResult

  mountMaybe :: a -> Bool -> MIO ()
  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 -> MIO 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 -> MIO Bool
  isMounted dev = mountedState <$> mountState dev

  mountState :: a -> MIO 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 -> MIO (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
  => MIO [a]
  -> MIO [(Header, ProtoAction)]
mountableToAction ms = mapM mkAction =<< ms

type MIO 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) (MIO ())

--------------------------------------------------------------------------------
-- 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' $
      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
  :: (HasLogFunc c, MonadReader c m, MonadUnliftIO m)
  => FilePath
  -> Maybe PasswordConfig
  -> T.Text
  -> m MountResult
mountSSHFS mountpoint pwdConfig remote =
  withPasswordGetter pwdConfig (run ["-o", "password_stdin"]) $ run [] ""
  where
    run other = runMount "sshfs" (other ++ [remote, T.pack mountpoint])

mountCIFS
  :: (HasLogFunc c, MonadReader c m, MonadUnliftIO m)
  => Bool
  -> T.Text
  -> FilePath
  -> Maybe CIFSOpts
  -> Maybe PasswordConfig
  -> m 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
  :: (HasLogFunc c, MonadReader c m, MonadUnliftIO m)
  => FilePath
  -> Maybe PasswordConfig
  -> T.Text
  -> m 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 :: MonadIO m => [T.Text] -> T.Text -> m MountResult
runVeraCrypt args = runMount "sudo" (defaultArgs ++ args)
  where
    defaultArgs = ["/usr/bin/veracrypt", "--text", "--non-interactive"]

veracryptMountState :: MountConfig -> MIO 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 -> MIO FilePath
getAbsMountpoint MountConfig {mpPath = m} =
  asks $ flip appendRoot (T.unpack m) . mountconfVolatilePath

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

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

type PasswordGetter m = m (Maybe T.Text)

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

runBitwarden :: (HasLogFunc c, MonadReader c m, MonadUnliftIO m) => T.Text -> PasswordGetter m
runBitwarden pname =
  ((password . login) <=< L.find (\i -> name i == pname))
    <$> getItems

runPromptLoop :: MonadUnliftIO m => Natural -> PasswordGetter m -> PasswordGetter m
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 :: (HasLogFunc c, MonadReader c m, MonadUnliftIO m) => PasswordConfig -> PasswordGetter m
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
  :: (HasLogFunc c, MonadReader c m, MonadUnliftIO m)
  => Maybe PasswordConfig
  -> (T.Text -> m MountResult)
  -> m MountResult
  -> m 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 :: MIO [(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 :: MIO [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 :: MIO [(Header, ProtoAction)]
getMTPActions = mountableToAction getMTPDevices

mtpExeInstalled :: MonadIO m => m Bool
mtpExeInstalled = isJust <$> findExecutable mtpExe

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

  groupHeader _ = MTPFSHeader

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

notifyMountResult :: MonadIO m => Bool -> T.Text -> MountResult -> m ()
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

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

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

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

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

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

runMountSudoMaybe' :: MonadIO m => Bool -> T.Text -> [T.Text] -> [(T.Text, T.Text)] -> m 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' :: MonadIO m => T.Text -> T.Text -> [T.Text] -> [(T.Text, T.Text)] -> m 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 -> MIO MountResult -> MIO 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 -> MIO MountResult -> MIO MountResult
withTmpMountDir m =
  rmDirOnMountError m
    . bracketOnError_ (mkDirMaybe m) (rmDirMaybe m)

-- | Run an unmount command and remove the mountpoint if no errors occur
runAndRemoveDir :: FilePath -> MIO MountResult -> MIO 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 -> MIO ()
mkDirMaybe fp = whenInMountDir fp $ io $ createDirectoryIfMissing True fp

rmDirMaybe :: FilePath -> MIO ()
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 -> MIO () -> MIO ()
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