{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE FlexibleInstances #-} -------------------------------------------------------------------------------- -- | 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 Control.Monad import Control.Monad.Reader import Data.Either import Data.List import Data.List.Split (splitOn) import qualified Data.Map as M import Data.Maybe import qualified Data.Text as T import qualified Data.Vector as V import Data.Yaml import Rofi.Command import Text.Printf import System.Console.GetOpt import System.Directory import System.Environment import System.FilePath.Posix import System.Posix.User (getEffectiveUserName) import System.Process import UnliftIO.Exception main :: IO () main = getArgs >>= parse parse :: [String] -> IO () parse args = case getOpt Permute options args of (o, n, []) -> runMounts $ foldl (flip id) (defaultOpts n) o (_, _, errs) -> ioError $ userError $ 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 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 String , optsUnmount :: Bool , optsRofiArgs :: [String] } 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 $ _staticconfigTmpPath =<< static let staticDevs = maybe M.empty _staticconfigDevices static let verbose = fromMaybe False $ _staticconfigVerbose =<< static let mountconf = MountConf { mountconfVolatilePath = tmpPath , mountconfRofiArgs = optsRofiArgs opts , mountconfStaticDevs = staticDevs , mountconfVerbose = verbose } let byAlias = mountByAlias $ optsUnmount opts let byPrompt = runPrompt =<< getGroups runRofiIO mountconf $ maybe byPrompt byAlias $ optsAlias opts parseStaticConfig :: FilePath -> IO (Maybe StaticConfig) parseStaticConfig p = do res <- decodeFileEither p case res of Left e -> print e >> return Nothing Right c -> return $ Just c runPrompt :: RofiConf c => [RofiGroup c] -> RofiIO c () runPrompt gs = selectAction $ emptyMenu { groups = gs , prompt = Just "Select Device" } getGroups :: RofiMountIO [RofiGroup MountConf] getGroups = do staticDevs <- asks mountconfStaticDevs staticActions <- mapM mkAction $ configToTree' staticDevs removableActions <- mapM mkAction =<< getRemovableDevices mtpActions <- mapM mkAction =<< getMTPDevices return $ mapMaybe mkGroup $ groupBy (\(hx, _) (hy, _) -> hx == hy) $ sortBy (\(hx, _) (hy, _) -> compare hx hy) $ staticActions ++ removableActions ++ mtpActions mountByAlias :: Bool -> String -> RofiMountIO () mountByAlias unmountFlag alias = do static <- asks mountconfStaticDevs mapM_ (`mountMaybe` unmountFlag) $ configToTree static <$> M.lookup alias static mkGroup :: [(Header, ProtoAction [String])] -> Maybe (RofiGroup MountConf) mkGroup [] = Nothing mkGroup as = let ((Header title _, _):_) = as in Just $ titledGroup title $ toRofiActions $ alignEntries $ fmap snd as alignSep :: String alignSep = " | " alignEntries :: [ProtoAction [String]] -> [(String, RofiMountIO ())] alignEntries ps = zip (align es) as where (es, as) = unzip $ fmap (\(ProtoAction e a) -> (e, a)) ps align = fmap (intercalate alignSep) . transpose . mapToLast pad . transpose pad xs = let m = getMax xs in fmap (\x -> take m (x ++ repeat ' ')) xs getMax = maximum . fmap length mapToLast _ [] = [] mapToLast _ [x] = [x] mapToLast f (x:xs) = f x : mapToLast f xs -------------------------------------------------------------------------------- -- | Global config used in the reader monad stack data MountConf = MountConf { mountconfVolatilePath :: FilePath , mountconfRofiArgs :: [String] , mountconfStaticDevs :: M.Map String TreeConfig , mountconfVerbose :: Bool } deriving Show instance RofiConf MountConf where defArgs MountConf { mountconfRofiArgs = a } = a -------------------------------------------------------------------------------- -- | Mountable typeclass -- -- Class to provide common interface for anything that can be mounted. class Mountable a where -- | Mount the given type (or dismount if False is passed) mount :: a -> Bool -> RofiMountIO () mountMaybe :: a -> Bool -> RofiMountIO () mountMaybe dev mountFlag = do mounted <- isMounted dev verbose <- asks mountconfVerbose if mountFlag == mounted then mount dev mountFlag else when verbose notify' where notify' = io $ notify "dialog-information-symbolic" $ getLabel dev ++ " already mounted" -- | Check if the mounting utilities are present allInstalled :: a -> RofiMountIO Bool -- | Return a string representing the label of the device getLabel :: a -> String -- | Determine if the given type is mounted or not isMounted :: a -> RofiMountIO Bool -------------------------------------------------------------------------------- -- | 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 -> [String] 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 [String]) mkAction dev = do m <- isMounted dev i <- allInstalled dev let h = groupHeader dev let action = when i $ mountMaybe dev m let entry = case fmtEntry dev of (e:es) -> (mountedPrefix m i ++ e):es _ -> [] return (h, ProtoAction entry action) where mountedPrefix False True = " " mountedPrefix True True = "* " mountedPrefix _ False = "! " type RofiMountIO a = RofiIO MountConf a data Header = Header String Integer deriving (Show, Eq) data ProtoAction a = ProtoAction a (RofiMountIO ()) instance Ord Header where compare (Header _ x) (Header _ y) = compare x y -------------------------------------------------------------------------------- -- | Static device configuration -- -- Static devices are defined in a YAML file. These types/instances describe how -- to parse said YAML file. defaultTries :: Integer defaultTries = 2 (.:&) :: FromJSON a => Object -> T.Text -> Parser (V.Vector a) (.:&) o t = o .:? t .!= V.empty data MountConfig = MountConfig { _mountMountPoint :: FilePath , _mountLabel :: Maybe String } deriving Show instance FromJSON MountConfig where parseJSON = withObject "devices" $ \o -> MountConfig <$> o .: "mountpoint" <*> o .:? "label" data BitwardenConfig = BitwardenConfig { _bitwardenKey :: String , _bitwardenTries :: Integer } deriving Show instance FromJSON BitwardenConfig where parseJSON = withObject "bitwarden" $ \o -> BitwardenConfig <$> o .: "key" <*> o .:? "tries" .!= defaultTries newtype LibSecretConfig = LibSecretConfig { _libsecretAttributes :: M.Map String String } deriving Show instance FromJSON LibSecretConfig where parseJSON = withObject "libsecret" $ \o -> LibSecretConfig <$> o .: "attributes" newtype PromptConfig = PromptConfig { _promptTries :: Integer } deriving Show instance FromJSON PromptConfig where parseJSON = withObject "libsecret" $ \o -> PromptConfig <$> o .: "tries" .!= defaultTries data PasswordConfig = PasswordConfig { _passwordBitwarden :: Maybe BitwardenConfig , _passwordLibSecret :: Maybe LibSecretConfig , _passwordPrompt :: Maybe PromptConfig } deriving Show instance FromJSON PasswordConfig where parseJSON = withObject "password" $ \o -> PasswordConfig <$> o .:? "bitwarden" <*> o .:? "libsecret" <*> o .:? "prompt" data DataConfig = VeracryptConfig { _veracryptVolume :: String , _veracryptPassword :: Maybe PasswordConfig } | SSHFSConfig { _sshfsRemote :: String } | CIFSConfig { _cifsRemote :: String , _cifsPassword :: Maybe PasswordConfig } deriving Show data DeviceConfig = DeviceConfig { _deviceMount :: MountConfig , _deviceData :: DataConfig } deriving Show data TreeConfig = TreeConfig { _treeParent :: DeviceConfig , _treeChildren :: V.Vector String } deriving Show instance FromJSON TreeConfig where parseJSON = withObject "devices" $ \o -> do devType <- o .: "type" deps <- o .:& "depends" mountconf <- o .: "mount" devData <- case (devType :: String) of "cifs" -> CIFSConfig <$> o .: "remote" <*> o .:? "password" "sshfs" -> SSHFSConfig <$> o .: "remote" "veracrypt" -> VeracryptConfig <$> o .: "volume" <*> o .:? "password" -- TODO make this skip adding an entry to the map rather than -- skipping the map entirely _ -> fail $ "unknown device type: " ++ devType return $ TreeConfig { _treeParent = DeviceConfig { _deviceMount = mountconf , _deviceData = devData } , _treeChildren = deps } data StaticConfig = StaticConfig { _staticconfigTmpPath :: Maybe String , _staticconfigVerbose :: Maybe Bool , _staticconfigDevices :: M.Map String TreeConfig } deriving Show instance FromJSON StaticConfig where parseJSON = withObject "devices" $ \o -> StaticConfig <$> o .:? "mountdir" <*> o .:? "verbose" <*> o .: "devices" -------------------------------------------------------------------------------- -- | 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 isMounted (Tree p _) = isMounted 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{ _cifsRemote = r } = r target SSHFSConfig{ _sshfsRemote = r } = r target VeracryptConfig{ _veracryptVolume = v } = v groupHeader (Tree DeviceConfig{ _deviceData = d } _) = case d of CIFSConfig{} -> Header "CIFS Devices" 0 SSHFSConfig{} -> Header "SSHFS Devices" 1 VeracryptConfig{} -> Header "Veracrypt Devices" 2 configToTree' :: M.Map String TreeConfig -> [StaticConfigTree] configToTree' devMap = configToTree devMap <$> M.elems devMap configToTree :: M.Map String TreeConfig -> TreeConfig -> StaticConfigTree configToTree devMap TreeConfig{ _treeParent = p, _treeChildren = c } = Tree p $ fmap go V.toList 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 c@DeviceConfig{ _deviceMount = m, _deviceData = devData} False = do m' <- getAbsMountpoint m bracketOnError_ (mkDirMaybe m') (rmDirMaybe m') $ mount' m' where mount' mountpoint = io $ case devData of SSHFSConfig{ _sshfsRemote = r } -> do runMountNotify "sshfs" [r, mountpoint] (getLabel c) False CIFSConfig{ _cifsPassword = p } -> do res <- case p of Just pwd -> do pwd' <- maybe [] (\p' -> [("PASSWD", p')]) <$> configToPwd pwd readCmdEither' "mount" [mountpoint] "" pwd' Nothing -> readCmdEither "mount" [mountpoint] "" notifyMounted (isRight res) False (getLabel c) VeracryptConfig{ _veracryptPassword = getPwd, _veracryptVolume = v } -> maybe (runVeraCryptWith "" []) (runVeraCryptWithPwd =<<) (configToPwd <$> getPwd) where label = getLabel c runVeraCryptWithPwd = maybe notifyFail (\p -> runVeraCryptWith p ["--stdin"]) runVeraCryptWith stdin args = (\res -> notifyMounted (isRight res) False label) =<< runVeraCrypt stdin ([v, mountpoint] ++ args) notifyFail = notify "dialog-error-symbolic" $ printf "Failed to get volume password for %s" label mount c@DeviceConfig{ _deviceMount = m, _deviceData = VeracryptConfig{} } True = do m' <- getAbsMountpoint m res <- io $ runVeraCrypt "" ["-d", m'] io $ notifyMounted (isRight res) True (getLabel c) mount c@DeviceConfig{ _deviceMount = m } True = umountNotify (getLabel c) =<< getAbsMountpoint m allInstalled DeviceConfig{ _deviceData = devData } = io $ isJust <$> findExecutable (exe devData) where exe SSHFSConfig{} = "sshfs" exe CIFSConfig{} = "mount.cifs" exe VeracryptConfig{} = "veracrypt" isMounted DeviceConfig{ _deviceMount = m } = (io . isDirMounted) =<< getAbsMountpoint m getLabel DeviceConfig { _deviceMount = MountConfig { _mountMountPoint = p, _mountLabel = l } } = fromMaybe (takeFileName p) l getAbsMountpoint :: MountConfig -> RofiMountIO FilePath getAbsMountpoint MountConfig{ _mountMountPoint = m } = asks $ flip appendRoot m . mountconfVolatilePath -- NOTE: the user is assumed to have added themselves to the sudoers file so -- that this command will work runVeraCrypt :: String -> [String] -> IO (Either (Int, String, String) String) runVeraCrypt stdin args = do readCmdEither "sudo" (defaultArgs ++ args) stdin where defaultArgs = ["/usr/bin/veracrypt", "--text", "--non-interactive"] -------------------------------------------------------------------------------- -- | Password-getting functions for static devices type PasswordGetter = IO (Maybe String) runSecret :: [(String, String)] -> PasswordGetter runSecret kvs = readCmdSuccess "secret-tool" ("lookup":kvs') "" where kvs' = concatMap (\(k, v) -> [k, v]) kvs runBitwarden :: String -> PasswordGetter runBitwarden pname = ((password . login) <=< find (\i -> name i == pname)) <$> getItems runPromptLoop :: Integer -> 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{ _bitwardenKey = k, _bitwardenTries = n }) = runPromptLoop n $ runBitwarden k getBW _ = return Nothing getLS = maybe (return Nothing) (runSecret . M.toList . _libsecretAttributes) getPrompt = maybe (return Nothing) (flip runPromptLoop readPassword . _promptTries) runMaybe x y = (\r -> if isNothing r then y else return r) =<< x -------------------------------------------------------------------------------- -- | 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 { deviceSpec :: String , label :: String } deriving (Eq, Show) instance Mountable Removable where mount Removable { deviceSpec = d, label = l } m = io $ runMountNotify "udisksctl" [c, "-b", d] l m where c = if m then "unmount" else "mount" allInstalled _ = fmap isJust $ io $ findExecutable "udisksctl" isMounted Removable { deviceSpec = d } = elem d <$> io curDeviceSpecs getLabel Removable { label = l } = l instance Actionable Removable where fmtEntry Removable { deviceSpec = d, label = l } = [l, d] groupHeader _ = Header "Removable Devices" 3 -- | 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 :: RofiConf c => RofiIO c [Removable] getRemovableDevices = fromLines toDev . lines <$> 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 splitBy ' ' line of ("":_) -> Nothing [_, "1", d, "", s] -> mk d $ s ++ " Volume" [_, "1", d, l, _] -> mk d l _ -> Nothing mk d l = Just $ Removable { deviceSpec = d, label = l } -------------------------------------------------------------------------------- -- | MTP devices data MTPFS = MTPFS { bus :: String , device :: String , mountpoint :: FilePath , description :: String } deriving (Eq, Show) instance Mountable MTPFS where mount MTPFS {..} False = do -- TODO add autodismount to options let dev = "-device=" ++ bus ++ "," ++ device bracketOnError_ (mkDirMaybe mountpoint) (rmDirMaybe mountpoint) (io $ runMountNotify "jmtpfs" [dev, mountpoint] description False) mount MTPFS { mountpoint = m, description = d } True = umountNotify d m -- | return True always since the list won't even show without jmtpfs allInstalled _ = return True isMounted MTPFS { mountpoint = dir } = io $ isDirMounted dir getLabel MTPFS { description = d } = d -- | Return list of all available MTP devices getMTPDevices :: RofiMountIO [MTPFS] getMTPDevices = do dir <- asks mountconfVolatilePath res <- io $ readProcess "jmtpfs" ["-l"] "" return $ fromLines (toDev dir) $ toDevList res where toDevList = reverse . takeWhile (not . isPrefixOf "Available devices") . reverse . lines toDev dir s = case splitOn ", " s of [busNum, devNum, _, _, desc, vendor] -> let d = unwords [vendor, desc] in Just $ MTPFS { bus = busNum , device = devNum , mountpoint = dir canonicalize d , description = d } _ -> Nothing canonicalize = mapMaybe repl repl c | c `elem` ("\"*/:<>?\\|" :: String) = Nothing | c == ' ' = Just '-' | otherwise = Just c instance Actionable MTPFS where fmtEntry d = [getLabel d] groupHeader _ = Header "MTP Devices" 5 -------------------------------------------------------------------------------- -- | Low-level mount functions -- ASSUME these will never fail because the format of /proc/mounts is fixed curMountField :: Int -> IO [String] curMountField i = fmap ((!! i) . words) . lines <$> readFile "/proc/mounts" curDeviceSpecs :: IO [String] curDeviceSpecs = curMountField 0 curMountpoints :: IO [String] curMountpoints = curMountField 1 -- 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) 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 () -> RofiIO MountConf () whenInMountDir fp f = do mDir <- asks mountconfVolatilePath when (mDir `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 runMountNotify :: String -> [String] -> String -> Bool -> IO () runMountNotify cmd args msg mounted = do res <- readCmdEither cmd args "" notifyMounted (isRight res) mounted msg umountNotify' :: String -> String -> FilePath -> RofiMountIO () umountNotify' cmd msg dir = finally (io $ runMountNotify cmd [dir] msg True) (rmDirMaybe dir) umountNotify :: String -> FilePath -> RofiMountIO () umountNotify = umountNotify' "umount" -- | Send a notification indicating the mount succeeded notifyMounted :: Bool -> Bool -> String -> IO () notifyMounted succeeded mounted label = notify icon body where (format, icon) = if succeeded then ("Successfully %sed %s", "dialog-information-symbolic") else ("Failed to %s %s", "dialog-error-symbolic") m = if mounted then "unmount" else "mount" :: String body = printf format m label notify :: String -> String -> IO () notify icon body = void $ spawnProcess "notify-send" ["-i", icon, body] -------------------------------------------------------------------------------- -- | Other functions fromLines :: (String -> Maybe a) -> [String] -> [a] fromLines f = mapMaybe (f . stripWS) -- TODO this exists somewhere... splitBy :: Char -> String -> [String] splitBy delimiter = 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