{-# LANGUAGE UnicodeSyntax, LambdaCase, ScopedTypeVariables, FlexibleContexts #-}
module Main where
import Control.Monad
import Control.Monad.State
import Control.Lens
import Text.Read (readMaybe)
import System.Directory
import System.Environment
import System.Exit
import System.Process
import Control.Exception
import Data.Maybe
import Data.List (isPrefixOf, sort)
import qualified DMenu
-- | List the files in a directory or the empty list in case of an exception.
listDirectoryDef
:: MonadIO m
=> FilePath
-> m [FilePath]
listDirectoryDef p =
liftIO $ fmap sort $
listDirectory p `catch` (\(_ :: SomeException) -> pure [])
-- | Return all files in the /dev/ directory.
getDevs
:: MonadIO m
=> m [FilePath]
getDevs =
liftIO $ listDirectoryDef "/dev"
-- | Run a process.
runProc
:: MonadIO m
=> String -- ^ Program name
-> [String] -- ^ Program arguments
-> String -- ^ Input piped into @stdin@
-> m (Either String String) -- ^ Either the @stderr@ output, if the process failed, or the @stdout@ output.
runProc prog args sIn =
liftIO $ do
(exitCode, sOut, sErr) ←
readCreateProcessWithExitCode (proc prog args) sIn
pure $ case exitCode of
ExitSuccess → Right sOut
ExitFailure _ → Left sErr
-- | Run a process.
runProcOr
:: MonadIO m
=> String -- ^ Program name
-> [String] -- ^ Program arguments
-> String -- ^ Input piped into @stdin@
-> String -- ^ Default value, returned if process exits with failure
-> m String -- ^ Either the default value, or the @stdout@ output of the process
runProcOr prog args sIn sDef =
either (const sDef) id <$> runProc prog args sIn
-- | Keep 'String's which have certain prefixes.
filterWithPrefixes
:: [String] -- ^ List of prefixes
-> [String] -- ^ List of 'String's to filter
-> [String] -- ^ 'String's which have one of the prefixes as prefix.
filterWithPrefixes = \case
[] -> id
ps -> filter $ \s → any (`isPrefixOf` s) ps
-- | Drop an amount of number substrings separated by whitespace from a string.
skipNumbersAndWS
:: Int -- ^ How many numbers to skip
-> String
-> String
skipNumbersAndWS =
go False
where
go _ 0 s = dropWhile (`elem` [' ', '\t']) $ s
go _ _ "" = ""
go b n (c:s) | c `elem` ['0'..'9'] = go True n s
| otherwise = go False (if b then n-1 else n) s
-- | Display a byte size as 'String'.
showBytes
:: Integer -- ^ Number of bytes for @1 KB@. Usually @1024@ or @1000@.
-> Integer -- ^ Number of bytes to display as 'String'.
-> String
showBytes k i
| i < k = show i ++ " B"
| i < k*k = show (i `div` k) ++ " KB"
| i < k*k*k = show (i `div` (k*k)) ++ " MB"
| i < k*k*k*k = show (i `div` (k*k*k)) ++ " GB"
| i < k*k*k*k*k = show (i `div` (k*k*k*k)) ++ " TB"
| otherwise = show (i `div` (k*k*k*k*k)) ++ " PB"
-- | Run @cat /proc/partitions@ and collect info about the capacities of
-- storage devices.
getPartitionInfos
:: MonadIO m
=> m [(FilePath, Integer)] -- ^ Map from device name to @String@ representation of it's capacity.
getPartitionInfos = do
sOut ← runProcOr "cat" ["/proc/partitions"] "" ""
forM (drop 2 $ lines sOut) $ \l → do
let numBlocks :: Integer = read $ head $ words $ skipNumbersAndWS 2 l
let devPath = head (words $ skipNumbersAndWS 3 l)
pure (devPath, numBlocks * 1024)
-- | Make all @String@s in a list have the same size, by filling missing
-- characters with spaces from the right.
fillWithSpR
:: [String]
-> [String]
fillWithSpR ss =
map f ss
where
f s = s ++ replicate (maxLength - length s) ' '
maxLength = maximum (map length ss)
-- | Make all @String@s in a list have the same size, by filling missing
-- characters with spaces from the left.
fillWithSpL
:: [String]
-> [String]
fillWithSpL ss =
map f ss
where
f s = replicate (maxLength - length s) ' ' ++ s
maxLength = maximum (map length ss)
-- | Spawn a @pmount@ process and retrieve a list of currently mounted devices.
getMountedDevs
:: MonadIO m
=> m [FilePath]
getMountedDevs =
liftIO $ do
(exitCode, sOut, _sErr) <- readCreateProcessWithExitCode
(proc "pmount" []) ""
pure $ case exitCode of
ExitSuccess → map ((!! 0) . words) $ drop 1 $ init $ lines sOut
ExitFailure _ → []
-- | Parse the command line arguments
readArgs
:: [String] -- ^ Arguments from 'getArgs'
-> IO (Bool, [String], Integer) -- ^ (1) Whether the unmount flag was set, (2) a list of device prefixes to filter for, (3) how much byte are 1 KB.
readArgs args =
execStateT (go $ words $ unwords args) (False, [], 1024)
where
go [] =
pure ()
go (a:as)
-- All arguments after "--" are passed to dmenu later.
| a == "--" =
pure ()
| a `elem` ["-f", "--filter"] = do
let (prefixes, as') = break ("-" `isPrefixOf`) as
_2 %= (++prefixes)
go as'
| a `elem` ["-u", "--unmount"] = do
_1 .= True
go as
| a `elem` ["-k", "--kilo"]
, a':as' <- as
, Just k <- readMaybe a'
, k > 0 = do
_3 .= k
go as'
| a `elem` ["-h", "--help"] =
liftIO $ do putStrLn usage; exitFailure
| a == "" =
go as
go _ =
liftIO $ do
putStrLn usage
exitFailure
main
:: IO ()
main = do
(unmount, prefixes, k) <- readArgs =<< getArgs
let (prog, devsM) | unmount = ("pumount", map (drop 5) . filterWithPrefixes (map ("/dev/"++) prefixes) <$> getMountedDevs)
| otherwise = ("pmount", filterWithPrefixes prefixes <$> getDevs)
devs ← devsM
devInfos ← map (_2 %~ showBytes k) <$> getPartitionInfos
let devInfos' = map (fromMaybe "" . flip lookup devInfos) devs
let devs' = zip devs $ zipWith (\x y → x++" "++y) (fillWithSpR devs) (fillWithSpL devInfos')
let dmenuOpts = do
DMenu.prompt .= drop 1 prog
DMenu.forwardExtraArgs
DMenu.selectWith dmenuOpts snd devs' >>= \case
Right (dev,_) | not (null dev) → callProcess prog [dev]
_ → pure ()
usage
:: String
usage =
unlines
[ "USAGE"
, " dmenu-pmount [OPTIONS] [-- DMENUOPTIONS]"
, ""
, " Mount or unmount a linux device with dmenu and pmount as user."
, ""
, " All arguments, after the first `--` argument, are directly passed to dmenu."
, ""
, "OPTIONS"
, " -u, --unmount"
, " Unmount devices displayed by `pmount` with `pumount`."
, " -f, --filter <DevPrefixList>"
, " Only display devices whose filenames have a certain prefix."
, " Example: `dmenu-pmount -f sd cdrom` only displays devices beginning with"
, " `sd` or `cdrom`, e.g. `/dev/sda2`."
, " -k, --kilo <Natural>"
, " How much byte are represented by 1 KB? Usually 1024 (default) or 1000."
, " -h, --help"
, " Display this message."
]