rofi-extras/app/rofi-devices.hs

502 lines
17 KiB
Haskell

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Main (main) where
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 qualified Data.Map.Ordered as O
import Data.Maybe
import Data.Text (unpack)
import Rofi.Command
import Text.Printf
import Text.Regex.TDFA
import Text.Wrap
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 = check >> getArgs >>= parse
type Password = IO (Maybe String)
data MountConf = MountConf
{ credentials :: M.Map String Password
, mountDir :: FilePath
, rofiArgs :: [String]
}
instance RofiConf MountConf where
defArgs MountConf { rofiArgs = a } = a
type DevicePasswords = M.Map String Password
initMountConf :: [String] -> IO MountConf
initMountConf a = conf <$> getEffectiveUserName
where
conf u = MountConf
{ credentials = M.empty
, mountDir = "/media" </> u
, rofiArgs = a
}
options :: [OptDescr (MountConf -> MountConf)]
options =
[ Option ['s'] ["secret"]
(ReqArg (\s m -> m { credentials = addGetSecret (credentials m) s } ) "SECRET")
$ wrap "Use libsecret to retrieve password for DIR using ATTR/VAL pairs. \
\The pairs will be supplied to a 'secret-tool lookup' call. \
\ Argument is formatted like 'DIR:ATTR1=VAL1,ATTR2=VAL2...'"
, Option ['d'] ["directory"]
(ReqArg (\s m -> m { mountDir = s } ) "DIR")
$ wrap "The DIR in which new mountpoints will be created. This is assumed \
\to be writable to the current user, and will be used for fuse \
\entries as well as user mounts in fstab. For the latter, it is \
\assumed that all user mounts contain this directory if a \
\mountpoint does not already exist for them. If not given this will \
\default to '/media/USER'."
, Option ['p'] ["password"]
(ReqArg (\s m -> m { credentials = addGetPrompt (credentials m) s } ) "DIR")
"Prompt for password when mounting DIR."
]
where
wrap = unpack . wrapText defaultWrapSettings 40
parse :: [String] -> IO ()
parse args = case getOpt Permute options args of
(o, n, []) -> initMountConf n >>= \i -> runMounts $ foldl (flip id) i o
(_, _, errs) -> ioError $ userError $ concat errs ++ usageInfo h options
where
h = "Usage: rofi-dev [OPTIONS] [-- ROFI-OPTIONS]"
addGetSecret :: DevicePasswords -> String -> DevicePasswords
addGetSecret pwds c = case splitPrefix c of
(dir, ":", r) -> addPasswordGetter pwds dir $ runGetSecret
$ mapMaybe (toCell . splitEq) $ splitBy ',' r
_ -> pwds
where
splitPrefix s = s =~ (":" :: String) :: (String, String, String)
splitEq e = e =~ ("=" :: String) :: (String, String, String)
toCell (k, "=", v) = Just (k, v)
toCell _ = Nothing
runGetSecret :: [(String, String)] -> Password
runGetSecret kvs = readCmdSuccess "secret-tool" ("lookup":kvs') ""
where
kvs' = concatMap (\(k, v) -> [k, v]) kvs
addGetPrompt :: DevicePasswords -> String -> DevicePasswords
addGetPrompt pwds dir = addPasswordGetter pwds dir readPassword
addPasswordGetter :: DevicePasswords -> String -> IO (Maybe String) -> DevicePasswords
addPasswordGetter pwds key f = M.insert key f pwds
-- runGetBwPwd :: [(String, String)] -> IO (Maybe String)
-- runGetBwPwd = undefined
-- getPassword :: Credentials -> IO (Maybe String)
-- getPassword NoCredentials = return Nothing
-- getPassword (Secret kvs) = do
-- let kvs' = concat [[a, b] | (a, b) <- M.toList kvs]
-- readCmdSuccess "secret-tool" ("lookup":kvs') ""
-- TODO
check :: IO ()
check = return ()
runMounts :: MountConf -> IO ()
runMounts c = runRofiIO c $ runPrompt =<< getGroups
runPrompt :: RofiConf c => [RofiGroup c] -> RofiIO c ()
runPrompt gs = selectAction $ emptyMenu
{ groups = gs
, prompt = Just "Select Device"
}
getGroups :: RofiIO MountConf [RofiGroup MountConf]
getGroups = do
fstab <- readFSTab
sequence
[ mkGroup "SSHFS Devices" $ sshfsDevices fstab
, mkGroup "CIFS Devices" $ cifsDevices fstab
, mkGroup "Removable Devices" =<< getRemovableDevices
, mkGroup "MTP Devices" =<< getMTPDevices
]
mkGroup :: Mountable d => String -> [d] -> RofiIO MountConf (RofiGroup MountConf)
mkGroup header devs = titledGroup header . alignEntries . toRofiActions
<$> mapM mkAction devs
alignSep :: String
alignSep = " | "
alignSepPre :: String
alignSepPre = "@@@"
alignEntries :: RofiActions c -> RofiActions c
alignEntries = O.fromList . withKeys . O.assocs
where
withKeys as = let (ks, vs) = unzip as in zip (align ks) vs
align ks = fmap (intercalate alignSep)
$ transpose
$ mapToLast pad
$ transpose
$ fmap (splitOn alignSepPre) ks
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
-- | Class and methods for type representing mountable devices
class Mountable a where
-- | Mount the given type (or dismount if False is passed)
mount :: a -> Bool -> RofiIO MountConf ()
-- | Check if the mounting utilities are present
allInstalled :: a -> RofiIO MountConf Bool
-- | Return a string to go in the Rofi menu for the given type
fmtEntry :: a -> String
-- | Determine if the given type is mounted or not
isMounted :: a -> RofiIO MountConf Bool
-- | 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 -> RofiIO MountConf (String, RofiIO MountConf ())
mkAction dev = do
m <- isMounted dev
i <- allInstalled dev
let a = when i $ mount dev m
let s = mountedPrefix m i ++ fmtEntry dev
return (s, a)
where
mountedPrefix False True = " "
mountedPrefix True True = "* "
mountedPrefix _ False = "! "
-- | Key/val pairs to represent mount options. A Nothing for the value signifies
-- a standalone option (eg 'rw' and 'ro')
type MountOptions = M.Map String (Maybe String)
-- | Given a string of comma-separated 'key=val' pairs, return a mount options
-- map
parseMountOptions :: String -> MountOptions
parseMountOptions s = M.fromList $ toCell . splitEq <$> splitBy ',' s
where
splitEq e = e =~ ("=" :: String) :: (String, String, String)
toCell (k, "=", v) = (k, Just v)
toCell (k, _, _) = (k, Nothing)
-- -- | Given a mount options map, return a string of comma separated items
-- fmtMountOptions :: MountOptions -> String
-- fmtMountOptions = intercalate "," . fmap fromCell . M.toList
-- where
-- fromCell (k, Just v) = k ++ "=" ++ v
-- fromCell (k, Nothing) = k
--------------------------------------------------------------------------------
-- | Removable devices
--
-- A device which can be removed (which is all the devices we care about)
-- This can be minimally described by a device DEVICESPEC and LABEL.
data Removable = Removable
{ deviceSpec :: String
, label :: String
}
deriving (Eq, Show)
instance Mountable Removable where
-- | (Un)mount the device using udiskctl
mount Removable { deviceSpec = d, label = l } m = io $ do
res <- readCmdEither "udisksctl" [cmd, "-b", d] ""
notifyMounted (isRight res) m l
where
cmd = if m then "unmount" else "mount"
-- | Need udisksctl to mount and umount
allInstalled _ = fmap isJust $ io $ findExecutable "udisksctl"
-- | Check if the device is mounted using /proc/mount
isMounted Removable { deviceSpec = d } = elem d <$> io curDeviceSpecs
-- | Format the Rofi entry like 'LABEL - PATH' and add a star in the front
-- if the device is mounted
fmtEntry Removable { deviceSpec = d, label = l } = l ++ alignSepPre ++ d
-- | 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 = mapMaybe toDev
. lines
. stripWS
<$> 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 }
--------------------------------------------------------------------------------
-- | CIFS Devices
--
-- This wraps the Removable device (since it is removable) and also adds its
-- own mount options and credentials for authentication.
data CIFS = CIFS Removable FilePath (Maybe Password)
instance Mountable CIFS where
mount (CIFS Removable{ label = l } m getPwd) False =
bracketOnError_
(mkDirMaybe m)
(rmDirMaybe m)
$ io $ do
res <- case getPwd of
Just pwd -> do
p <- maybe [] (\p -> [("PASSWD", p)]) <$> pwd
readCmdEither' "mount" [m] "" p
Nothing -> readCmdEither "mount" [m] ""
notifyMounted (isRight res) False l
mount (CIFS Removable{ label = l } m _) True = umountNotify l m
allInstalled _ = io $ isJust <$> findExecutable "mount.cifs"
isMounted (CIFS _ dir _) = io $ isDirMounted dir
fmtEntry (CIFS r _ _) = fmtEntry r
-- TODO this smells like something that should be in a typeclass
fstabToCIFS :: FSTabEntry -> RofiIO MountConf CIFS
fstabToCIFS FSTabEntry{ fstabSpec = s, fstabDir = d, fstabOptions = o } = do
-- If the options specify "guest" don't require a password. Else try to find a
-- means to get the password from the command line options provided for the
-- this mountpoint. If nothing is found, prompt for a password. In any case,
-- the output will be passed to env variable PASSWD when mounting this cifs
-- directory and cause it to fail. Setting the env variable is necessary as
-- the cifs mount call will prompt for a password and hang otherwise.
pwd <- if M.member "guest" o
then return Nothing
else Just . M.findWithDefault readPassword d <$> asks credentials
let r = Removable { deviceSpec = smartSlashPrefix s, label = takeFileName d }
return $ CIFS r d pwd
where
smartSlashPrefix a = if "//" `isPrefixOf` a then a else "//" ++ a
--------------------------------------------------------------------------------
-- | SSHFS Devices
--
-- This wraps the Removable device (since it is removable) and also adds its
-- own mount options. If the path does not point to an aliased entry in the ssh
-- config that specifies the port, hostname, user, and identity file, these
-- need to be passed as mount options.
data SSHFS = SSHFS Removable FilePath
instance Mountable SSHFS where
mount (SSHFS Removable{ label = l } m) False =
bracketOnError_
(mkDirMaybe m)
(rmDirMaybe m)
$ io $ do
res <- readCmdEither "mount" [m] ""
notifyMounted (isRight res) False l
mount (SSHFS Removable{ label = l } m) True = umountNotify l m
allInstalled _ = fmap isJust $ io $ findExecutable "sshfs"
isMounted (SSHFS _ dir) = io $ isDirMounted dir
fmtEntry (SSHFS r _) = fmtEntry r
fstabToSSHFS :: FSTabEntry -> RofiIO MountConf SSHFS
fstabToSSHFS FSTabEntry{ fstabSpec = s, fstabDir = d } = return $ SSHFS r d
where
r = Removable { deviceSpec = s, label = takeFileName d }
--------------------------------------------------------------------------------
-- | 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 $ do
res <- readCmdEither "jmtpfs" [dev, mountpoint] ""
notifyMounted (isRight res) False description
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
fmtEntry MTPFS { description = d } = d
getMTPDevices :: RofiIO MountConf [MTPFS]
getMTPDevices = do
dir <- asks mountDir
res <- io $ readProcess "jmtpfs" ["-l"] ""
return $ mapMaybe (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
-- TODO add truecrypt volumes (see tcplay, will need root)
--------------------------------------------------------------------------------
-- | Fstab devices
data FSTab = FSTab
{ sshfsDevices :: [SSHFS]
, cifsDevices :: [CIFS]
}
data FSTabEntry = FSTabEntry
{ fstabSpec :: String
, fstabDir :: FilePath
, fstabType :: String
, fstabOptions :: MountOptions
}
readFSTab :: RofiIO MountConf FSTab
readFSTab = do
let i = FSTab { sshfsDevices = [], cifsDevices = [] }
fstab <- io $ readFile "/etc/fstab"
foldM addFstabDevice i $ mapMaybe toEntry $ lines fstab
where
toEntry line = case words $ stripWS line of
(('#':_):_) -> Nothing
[spec, dir, fsType, opts, _, _] -> Just $ FSTabEntry
{ fstabSpec = spec
, fstabDir = dir
, fstabType = fsType
, fstabOptions = parseMountOptions opts
}
_ -> Nothing
addFstabDevice :: FSTab -> FSTabEntry -> RofiIO MountConf FSTab
addFstabDevice f@FSTab{..} e@FSTabEntry{..}
| M.notMember "users" fstabOptions = return f
| fstabType == "cifs" =
(\d -> f { cifsDevices = cifsDevices ++ [d] }) <$> fstabToCIFS e
| fstabType == "fuse.sshfs" =
(\d -> f { sshfsDevices = sshfsDevices ++ [d] }) <$> fstabToSSHFS e
| otherwise = return f
--------------------------------------------------------------------------------
-- | Low-level mount functions
-- ASSUME these will never fail because the format of this file 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
mkDirMaybe :: FilePath -> RofiIO MountConf ()
mkDirMaybe fp = whenInMountDir fp $ io $ createDirectoryIfMissing True fp
rmDirMaybe :: FilePath -> RofiIO MountConf ()
rmDirMaybe fp = whenInMountDir fp $ unlessMountpoint fp
$ asks mountDir >>= io . rmUntil fp
where
rmUntil cur target = unless (target == cur) $ do
removePathForcibly cur
rmUntil (takeDirectory cur) target
whenInMountDir :: FilePath -> RofiIO MountConf () -> RofiIO MountConf ()
whenInMountDir fp f = do
mDir <- asks mountDir
when (mDir `isPrefixOf` fp) f
unlessMountpoint :: FilePath -> RofiIO MountConf () -> RofiIO MountConf ()
unlessMountpoint fp f = do
mounted <- io $ isDirMounted fp
unless mounted f
umountNotify :: String -> FilePath -> RofiIO MountConf ()
umountNotify label dir = finally cmd $ rmDirMaybe dir
where
cmd = io $ do
res <- readCmdEither "umount" [dir] ""
notifyMounted (isRight res) True label
isDirMounted :: FilePath -> IO Bool
isDirMounted fp = elem fp <$> curMountpoints
--------------------------------------------------------------------------------
-- | Other functions
-- 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
notifyMounted :: Bool -> Bool -> String -> IO ()
notifyMounted succeeded mounted label = void $ spawnProcess "notify-send" [msg]
where
f = if succeeded then "Successfully %sed %s" else "Failed to %s %s"
m = if mounted then "unmount" else "mount" :: String
msg = printf f m label