path: root/lib/PacketTranscoder.hs

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

import Debug.Trace
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 KeyRing.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


data PassphraseResponse = ObtainedPassphrase S.ByteString
                        | CanceledPassphrase
                        | NextPassphrase
 deriving Show

type PassphraseSource = (SymmetricAlgorithm,S2K) -> MappedPacket -> [IO PassphraseResponse]

interpretPassSpec :: InputFileContext
                        -> (Maybe MappedPacket, Map.Map KeyKey (OriginMapped Query))
                        -> PassphraseSpec
                        -> IO (KikiCondition (PassphraseSource, IO ()) )
interpretPassSpec ctx keys PassphraseSpec { passSpecPassFile = fd
                                          , passSpecKeySpec = keyspec
                                          , passSpecRingFile = inputfile } = do
    getpw <-
        cachedContents (Just "Enter possibly multi-line passphrase (Press CTRL-D when finished):\n")
                       ctx
                       fd
    let matchkey fp mp = matchpr fp (packet mp) == fp
        matchfile file mp = Map.member file (locations mp)
        specializers = [ fmap matchkey keyspec, fmap matchfile inputfile]
        specialize alg mp =
            if and $ map (\f -> f mp) $ catMaybes specializers
                then [ObtainedPassphrase <$> getpw]
                else []
    return $ KikiSuccess (specialize, return ())

interpretPassSpec ctx keys PassphraseAgent = do
    mb <- session
    fromMaybe (return AgentConnectionFailure) $ do
    agent <- mb
    Just $ do
    let cacheSearch alg mp =
            case getQueries keys mp of
                [] -> []
                kqry:qs -> -- trace ("queries="++show (kqry:qs)) $
                            map (sendQuery agent (AskNot,NextPassphrase)) (kqry:qs)
                            ++ sendQuery agent (initial_ask,CanceledPassphrase) kqry
                                : replicate 3 (sendQuery agent (AskAgain "Bad passphrase",CanceledPassphrase) kqry)
                 where
                   srcalg = symmetric_algorithm $ packet mp

                   initial_ask | Unencrypted <- srcalg = AskNew
                               | otherwise             = AskExisting

    return $ KikiSuccess (cacheSearch, quit agent)

sendQuery :: GnuPGAgent -> (QueryMode,PassphraseResponse) -> OriginMapped Query -> IO PassphraseResponse
sendQuery agent (ask,failure) qry = do
    mbpw <- getPassphrase agent ask (packet qry)
    case mbpw of
        Nothing -> do
            S8.hPutStrLn stderr $ S8.pack $ "Failed to get passphrase "++show failure
            return failure
        Just pw -> do
            -- S8.hPutStrLn stderr $ S8.pack $ "received pw: "++pw
            return $ ObtainedPassphrase $ S8.pack pw

getQueries :: (Maybe MappedPacket,Map KeyKey (OriginMapped Query)) -> OriginMapped Packet -> [OriginMapped Query]
getQueries (workingkey,keys) mp =
 let kk = keykey $ packet mp
 in case Map.lookup kk keys of
     Just qryk ->
         case queryMainKey (packet qryk) of
            Just maink ->
                let kkmain = keykey maink
                in case Map.lookup kkmain keys of
                    Just qrym -> qryk : qrym : (Map.elems . Map.delete kkmain . Map.delete kk) keys
                    Nothing -> [ qryk ]
            Nothing -> [ qryk ]
     Nothing ->
        -- This is probably a newly imported key.  We'll treat the current working key as it's main key.
        -- trace ("getQueries cache miss "++show (fingerprint $ packet mp)) []
        let (qryk,qrym)
             = fromMaybe (Query (packet mp) "anonymous2" Nothing,Nothing) $ do
                    guard $ is_subkey (packet mp)
                    working <- fmap packet workingkey
                    q <- Map.lookup (keykey working) keys
                    return ( Query (packet mp) (queryUID $ packet q) (Just working), Just q)
            delm = case workingkey of
                      Nothing -> id
                      Just wk -> Map.delete (keykey $ packet wk)
        in [ fmap (const qryk) mp ] ++ maybeToList qrym ++ (Map.elems . delm . Map.delete kk) keys


makeTranscoder :: IORef (Map.Map (KeyKey,SymmetricAlgorithm,S2K) Packet)
                -> IO (KikiCondition (PassphraseSource, IO ()))
                -> ( (SymmetricAlgorithm,S2K) -> MappedPacket -> IO (KikiCondition Packet) )
                -> ( (SymmetricAlgorithm,S2K) -> MappedPacket -> IO (KikiCondition Packet) )
makeTranscoder unkeysRef source next alg mp = do
    ssr <- source
    case ssr of
        KikiSuccess (getqs, quit) -> tries (quit >> next alg mp) $ getqs alg mp
        er                        -> return $ fmap (error "makeTranscoder") er
 where
    tries fin [] = fin
    tries fin (src:srcs) = do
        S8.hPutStrLn stderr "trying..."
        resp <- src
        case resp of
            CanceledPassphrase    -> return OperationCanceled
            NextPassphrase        -> tries fin srcs
            ObtainedPassphrase pw -> do
                let wkun = fromMaybe (packet mp) $ do
                                guard $ symmetric_algorithm (packet mp) /= Unencrypted
                                decryptSecretKey pw (packet mp)
                    kk = keykey (packet mp)
                    retryOrFail = tries fin srcs
                case symmetric_algorithm wkun of
                    Unencrypted -> do
                        modifyIORef unkeysRef (Map.insert (kk,Unencrypted,S2K 100 "") wkun)
                        ek <- case fst alg of
                                Unencrypted -> do
                                    S8.hPutStrLn stderr "decrypted packet"
                                    return $ Just wkun
                                _           -> encryptSecretKey pw (snd alg) (fst alg) wkun
                        case ek of
                            Nothing   -> do
                                S8.hPutStrLn stderr "failed to encrypt"
                                retryOrFail
                            Just wken -> do
                                S8.hPutStrLn stderr "success encrypted"
                                let (a,s) = alg
                                modifyIORef unkeysRef (Map.insert (kk,a,s) wken)
                                return $ KikiSuccess wken
                    _ -> do S8.hPutStrLn stderr "failed to decrypt"
                            retryOrFail

normalizeAlgorithm :: (SymmetricAlgorithm,S2K) -> (SymmetricAlgorithm,S2K)
normalizeAlgorithm (Unencrypted,_) = (Unencrypted,S2K 100 "")
normalizeAlgorithm alg             = alg

transcodeWithCache :: IORef (Map (KeyKey, SymmetricAlgorithm, S2K) Packet)
                      -> ( (SymmetricAlgorithm,S2K) -> MappedPacket -> IO (KikiCondition Packet) )
                      -> ( (SymmetricAlgorithm,S2K) -> MappedPacket -> IO (KikiCondition Packet) )
transcodeWithCache unkeysRef miss alg0 mp@MappedPacket{ packet = wk } = do
        let alg = normalizeAlgorithm alg0
        unkeys <- readIORef unkeysRef
        calls <- currentCallStack
        putStrLn $ concat [fingerprint wk," ", show (symmetric_algorithm wk,s2k wk)," --> ",show alg]
        mapM_ putStrLn calls
        if (symmetric_algorithm wk,s2k wk) == alg
            then return (KikiSuccess wk)
            else maybe (miss alg mp) (return . KikiSuccess)
                    $ Map.lookup (keykey wk,fst alg, snd alg) unkeys

tryInOrder [] _ _        = return BadPassphrase
tryInOrder [f] alg mp    = f alg mp
tryInOrder (f:fs) alg mp = do
    r <- f alg mp
    case r of
        KikiSuccess _ -> return r
        e -> do
            S8.hPutStrLn stderr $ S8.pack ("got "++errorString e++", trying next")
            tryInOrder fs alg mp

-- The transcoder works on 'MappedPacket' instead of 'Packet' so that
-- file-specific passphrases can be utilized.
makeMemoizingDecrypter :: KeyRingOperation -> InputFileContext
                        -> (Maybe MappedPacket, Map.Map KeyKey (OriginMapped Query))
                        -> IO PacketTranscoder
makeMemoizingDecrypter operation ctx (workingkey,keys) = do
    unkeysRef <- newIORef (Map.empty :: Map.Map (KeyKey,SymmetricAlgorithm,S2K) Packet)
    return $ tryInOrder $ map passSpecTranscoder chains ++ [ trans unkeysRef ]
 where
    (chains,passpecs) = span isChain $ sort $ opPassphrases operation
                        where isChain (PassphraseMemoizer {}) = True
                              isChain _                       = False
    srcs = map (interpretPassSpec ctx (workingkey,keys)) passpecs

    trans unkeysRef = transcodeWithCache unkeysRef (foldr (makeTranscoder unkeysRef) (\_ _ -> return BadPassphrase) srcs)

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