path: root/lib/PacketTranscoder.hs

{-# LANGUAGE TupleSections #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternGuards #-}
module PacketTranscoder where

import GHC.Stack
import Control.Monad
import Data.IORef
import Data.List
import Data.Maybe
import Data.OpenPGP
import Data.OpenPGP.Util
import GnuPGAgent
import qualified Data.ByteString       as S
import qualified Data.ByteString.Char8 as S8
import           Data.Map as Map (Map)
import qualified Data.Map as Map
import qualified Data.Traversable      as Traversable
import System.IO              ( stderr)
import System.Posix.IO        ( fdToHandle )
import Text.Show.Pretty as PP ( ppShow )
import Types
import ControlMaybe (handleIO_)

-- | Merge two representations of the same key, prefering secret version
-- because they have more information.
mergeKeyPacket :: String -> MappedPacket -> MappedPacket -> MappedPacket
mergeKeyPacket what key p =
    key { packet = minimumBy (keyCompare what) [packet key,packet p]
        , locations = Map.union (locations key) (locations p)
        }

-- | Compare different versions if the same key pair.  Public versions
-- are considered greater.  If the two packets do not represent the same
-- key or the packets are not keys at all, an error will result that
-- includes the context provided as the first argument.
keyCompare :: String -> Packet -> Packet -> Ordering
keyCompare what (SecretKeyPacket {}) (PublicKeyPacket {}) = LT
keyCompare what (PublicKeyPacket {}) (SecretKeyPacket {}) = GT
keyCompare what a b | keykey a==keykey b = EQ
keyCompare what a b = error $ unlines ["Unable to merge "++what++":"
                              , if isKey a then fingerprint a else ""
                              , PP.ppShow a
                              , if isKey b then fingerprint b else ""
                              , PP.ppShow b
                              ]

resolveInputFile :: InputFileContext -> InputFile -> [FilePath]
resolveInputFile ctx = resolve
 where
    resolve HomeSec = return (homesecPath ctx)
    resolve HomePub = return (homepubPath ctx)
    resolve (ArgFile f) = return f
    resolve _ = []

resolveForReport :: Maybe InputFileContext -> InputFile -> FilePath
resolveForReport mctx (Pipe fdr fdw) = resolveForReport mctx (ArgFile str)
 where str = case (fdr,fdw) of
                (0,1) -> "-"
                _     -> "&pipe" ++ show (fdr,fdw)
resolveForReport mctx (FileDesc fd) = resolveForReport mctx (ArgFile str)
 where str = "&" ++ show fd
resolveForReport mctx f = concat $ resolveInputFile ctx f
 where ctx = fromMaybe (InputFileContext "&secret" "&public") mctx

readInputFileS :: InputFileContext -> InputFile -> IO S.ByteString
readInputFileS ctx (Pipe fd _)   = fdToHandle fd >>= S.hGetContents
readInputFileS ctx (FileDesc fd) = fdToHandle fd >>= S.hGetContents
readInputFileS ctx inp = do
    let fname = resolveInputFile ctx inp
    fmap S.concat $ mapM S.readFile fname



-- | Reads contents of an 'InputFile' or returns the cached content from a prior call.
-- An optional prompt is provided and will be printed on stdout only in the case that
-- the provided 'InputFile' is 'FileDesc' 0 (i.e. stdin).
cachedContents :: Maybe S.ByteString -> InputFileContext -> InputFile -> IO (IO S.ByteString)
cachedContents maybePrompt ctx fd = do
    ref <- newIORef Nothing
    return $ get maybePrompt ref fd
 where
    trimCR bs = fst $ S.spanEnd (\x -> x==10 || x==13) bs

    get maybePrompt ref fd = do
        pw <- readIORef ref
        flip (flip maybe return) pw $ do
            if fd == FileDesc 0 then case maybePrompt of
                                    Just prompt -> S.hPutStr stderr prompt
                                    Nothing -> return ()
                               else return ()
            pw <- fmap trimCR $ readInputFileS ctx fd
            writeIORef ref (Just pw)
            return pw



makeMemoizingDecrypter :: KeyRingOperation -> InputFileContext
                        -> (Maybe MappedPacket, Map.Map KeyKey (OriginMapped Query))
                        -> IO PacketTranscoder
makeMemoizingDecrypter operation ctx (workingkey,keys) = do
   if null chains then  do
    -- (*) Notice we do not pass ctx to resolveForReport.
    -- This is because the merge function does not currently use a context
    -- and the pws map keys must match the MappedPacket locations.
    -- TODO: Perhaps these should both be of type InputFile rather than
    -- FilePath?
    -- pws :: Map.Map FilePath (IO S.ByteString)
    {-
    -- This disabled code obtained password sources from StreamInfo records.
    pws <-
        Traversable.mapM (cachedContents ctx . fromJust . pwfile . typ)
                         (Map.mapKeys (resolveForReport Nothing) -- see note (*) note above
                            $ Map.filter (isJust . pwfile . typ) $ opFiles operation)
    -}
        let prompt = Just "Enter possibly multi-line passphrase (Press CTRL-D when finished):\n"
        -- List of file-specific password sources.
        pws2 <-
            Traversable.mapM (cachedContents prompt ctx)
             $ Map.fromList $ mapMaybe
                    (\spec -> (,passSpecPassFile spec) `fmap` do
                        guard $ isNothing $ passSpecKeySpec spec
                        passSpecRingFile spec)
                    passspecs
        -- List of general password sources.
        defpw <- do
            Traversable.mapM (cachedContents prompt ctx . passSpecPassFile)
             $ listToMaybe $ filter (\sp -> isNothing (passSpecRingFile sp)
                                            && isNothing (passSpecKeySpec sp))
                                        $ passspecs
        unkeysRef <- newIORef (Map.empty :: Map.Map (KeyKey,SymmetricAlgorithm,S2K) Packet)
        return $ doDecrypt unkeysRef ({- pws `Map.union` -} pws2) defpw (not $ null agentspec)
    else let PassphraseMemoizer f = head chains
         in return f
 where
    (chains,passspecs0) = partition isChain $ opPassphrases operation
                            where isChain (PassphraseMemoizer {}) = True
                                  isChain _                       = False
    (agentspec,passspecs) = partition isAgent passspecs0
                            where isAgent PassphraseAgent = True
                                  isAgent _               = False
    doDecrypt :: IORef (Map.Map (KeyKey,SymmetricAlgorithm,S2K) Packet)
                -> Map.Map FilePath (IO S.ByteString)
                -> Maybe (IO S.ByteString)
                -> Bool
                -> (SymmetricAlgorithm,S2K)
                -> MappedPacket
                -> IO (KikiCondition Packet)
    doDecrypt unkeysRef pws defpw agent_requested (dest_alg,dest_s2k) mp0 = do
        unkeys <- readIORef unkeysRef
        let (mp,qry) = fromMaybe (mp0,Query (packet mp0) "anonymous2" Nothing)
                        $ mplus (do k <- Map.lookup kk keys
                                    return (mergeKeyPacket "decrypt" mp0 (fmap queryPacket k), packet k))
                                (do guard $ is_subkey (packet mp0)
                                    working <- fmap packet workingkey
                                    q <- fmap packet $ Map.lookup (keykey working) keys
                                    return (mp0, Query (packet mp0) (queryUID q) (Just working)))

            dest_s2k' | dest_alg==Unencrypted = S2K 100 ""
                      | otherwise             = dest_s2k

            wk = packet mp0
            kk = keykey wk
            fs = Map.keys $ locations mp

            decryptIt [] = return BadPassphrase
            decryptIt (getpw:getpws) = tries 1 getpw (decryptIt getpws)
             where
              tries count getpw recurse = do
                -- TODO: This function should use mergeKeyPacket to
                --  combine the packet with it's unspilled version before
                --  attempting to decrypt it.  Note: We are uninterested
                --  in the 'locations' field, so this would effectively
                --  allow you to run 'decryptIt' on an unencrypted public key
                --  to obtain it's secret key.
                handleIO_ (decryptIt []) $ do
                (pw,wants_retry) <- getpw (count,qry)
                let wkun = fromMaybe wk $ do
                                guard $ symmetric_algorithm (packet mp) /= Unencrypted
                                decryptSecretKey pw (packet mp)

                    retryOrFail
                        | Just clear <- wants_retry = if count < 4
                                                        then tries (count+1) getpw recurse
                                                        else clear >> recurse
                        | otherwise                 = recurse

                case symmetric_algorithm wkun of

                    Unencrypted -> do
                        writeIORef unkeysRef (Map.insert (kk,Unencrypted,S2K 100 "") wkun unkeys)
                        ek <- case dest_alg of
                                Unencrypted -> return $ Just wkun
                                _           -> encryptSecretKey pw dest_s2k' dest_alg wkun

                        case ek of
                            Nothing   -> retryOrFail
                            Just wken -> do
                                modifyIORef unkeysRef (Map.insert (kk,dest_alg,dest_s2k') wken)
                                return $ KikiSuccess wken

                    _ -> retryOrFail

            getpws = (map (const . fmap (,Nothing)) $ mapMaybe (`Map.lookup` pws) fs ++ maybeToList defpw) ++ [ agentpw | agent_requested ]

            agentpw (count,qry) = do
                    s <- session
                    fromMaybe (return ("",Nothing)) $ do
                    s <- s
                    Just $ do
                    let (firsttime,maink) | Just k <- (queryMainKey qry) = (2,k)
                                          | otherwise                    = (1,error "bug in makeMemoizingDecrypter")

                        alg  = symmetric_algorithm (queryPacket qry)

                        ask | count<firsttime  = AskNot
                            | count>firsttime  = AskAgain "Bad passphrase"
                            | count==firsttime = initial_ask
                            where
                               initial_ask | Unencrypted <- alg = AskNew
                                           | otherwise          = AskExisting

                        actual_qry | count<firsttime = qry { queryPacket = maink, queryMainKey = Nothing }
                                   | otherwise       = qry

                    let clear | count > firsttime = clearPassphrase s (queryPacket qry)
                              | otherwise         = return ()
                    clear
                    let sanitizeQry qry = (fingerprint $ queryPacket qry, queryUID qry, fmap fingerprint $ queryMainKey qry)
                    -- putStrLn $ "(count,firsttime,ask,qry,actual_qry)="++show (count,firsttime,ask,sanitizeQry qry, sanitizeQry actual_qry)
                    mbpw <- getPassphrase s ask actual_qry
                    quit s
                    return ( maybe "" S8.pack mbpw, guard (ask /= AskNew) >> Just clear )

        calls <- currentCallStack
        putStrLn $ concat [fingerprint wk," ", show (symmetric_algorithm wk,s2k wk)," --> ",show (dest_alg,dest_s2k)]
        mapM_ putStrLn calls
        if symmetric_algorithm wk == dest_alg && s2k wk == dest_s2k'
            then return (KikiSuccess wk)
            else maybe (decryptIt getpws)
                       (return . KikiSuccess)
                    $ Map.lookup (kk,dest_alg,dest_s2k') unkeys

keyQueries :: Maybe String -> Map InputFile (StreamInfo,Message) -> (Maybe MappedPacket, Map KeyKey (OriginMapped Query))
keyQueries grip ringPackets = (mwk, fmap makeQuery keys)
 where
    makeQuery (maink,mp,us) = mp { packet = q }
        where q = Query { queryPacket = packet mp
                        , queryUID = concat $ take 1 $ Map.keys $ Map.union us (getUIDS maink)
                        , queryMainKey = if is_subkey (packet mp)
                                            then maink `mplus` fmap packet mwk
                                            else Nothing
                        }

    getUIDS maink = fromMaybe Map.empty $ do
        k <- maink
        (_,_,mus) <- Map.lookup (keykey k) keys
        return mus

    -- | mwk 
    -- first master key matching the provided grip
    -- (the m is for "MappedPacket", wk for working key)
    mwk :: Maybe MappedPacket
    mwk = listToMaybe $ do
            fp <- maybeToList grip
            let matchfp mp
                 | not (is_subkey p) && matchpr fp p == fp = Just mp
                 | otherwise                               = Nothing
                    where p = packet mp
            Map.elems $ Map.mapMaybe matchfp $ fmap (\(_,p,_) -> p) $ keys

    keys = Map.foldl slurpkeys Map.empty
                        $ Map.mapWithKey filterSecrets ringPackets
     where
      filterSecrets :: InputFile -> (a,Message) -> [[MappedPacket]]
      filterSecrets f (_,Message ps)  = keygroups
        -- filter (isSecretKey . packet) mps
        where
          mps = zipWith (mappedPacketWithHint fname) ps [1..]
          fname = resolveForReport Nothing f -- (Just ctx) f
          keygroups = dropWhile (not . isSecretKey . packet . head)
                      $ groupBy (const $ not . isSecretKey . packet) mps
      slurpkeys :: (Map KeyKey (Maybe Packet,MappedPacket,Map String [Packet]))
                    -> [[MappedPacket]]
                    -> (Map KeyKey (Maybe Packet, MappedPacket,Map String [Packet]))
      slurpkeys m pss = Map.unionWith combineKeyKey m m2
        where

            m2 :: Map.Map KeyKey (Maybe Packet, MappedPacket, (Map.Map String [Packet]))
            m2 = Map.fromList
                    $ drop 1
                    $ scanl' build ([],(Nothing,error "bug in PacketTranscoder (1)", error "bug in PacketTranscoder (2)")) pss
                where
                    build (_,(main0,_,_)) ps = (kk,(main,kp,uidmap ps'))
                     where
                        main | is_subkey (packet kp) = main0
                             | otherwise             = Just $ packet kp
                        (kpkt,ps') = splitAt 1 ps
                        kp = head kpkt
                        kk = keykey . packet $ kp
            combineKeyKey (master1,mp,um) (master2,mp2,um2) = (master1 `mplus` master2,mp,Map.unionWith (++) um um2)
            uidmap ps = um2
             where
                ugs = dropWhile (not . isUserID . packet .head) $ groupBy (const $ not . isUserID . packet) ps
                um2 = Map.fromList
                        $ map (\(MappedPacket (UserIDPacket s) _:sigs)->(s,takeWhile isSignaturePacket $ map packet sigs)) ugs