3 files changed, 165 insertions, 50 deletions

diff --git a/lib/GnuPGAgent.hs b/lib/GnuPGAgent.hs
index 8bffd1b..165fdf2 100644
--- a/lib/GnuPGAgent.hs
+++ b/lib/GnuPGAgent.hs
@@ -3,6 +3,8 @@
 {-# LANGUAGE PatternGuards #-}
 module GnuPGAgent
     ( session
+    , Query(..)
+    , QueryMode(..)
     , getPassphrase
     , clearPassphrase
     , quit ) where
@@ -44,6 +46,12 @@ session = do
             connect sock (SockAddrUnix (gpghome ++ "/S.gpg-agent"))
             agent <- socketToHandle sock ReadWriteMode
             hSetBuffering agent LineBuffering
+            lookupEnv "DISPLAY" >>= \case
+                Just display -> do hPutStrLn agent ("option putenv DISPLAY="++display)
+                                   _ <- hGetLine agent
+                                   return ()
+                Nothing      -> return ()
+            -- TODO: GPG_TTY
             return $ Just $ GnuPGAgent agent
         Nothing -> do
             hPutStrLn stderr "Unable to find home directory."
@@ -64,23 +72,42 @@ clearPassphrase agent key = do
     let cmd = "clear_passphrase "++fingerprint key
     hPutStrLn (agentHandle agent) cmd
 
-getPassphrase :: GnuPGAgent -> Bool -> Packet -> String -> Maybe Packet -> IO (Maybe String)
-getPassphrase agent ask key uid masterkey = do
-    let askopt = if ask then "" else "--no-ask "
-        (er,pr,desc) = prompts key uid masterkey
+data Query = Query
+    { queryPacket  :: Packet
+    , queryUID     :: String
+    , queryMainKey :: Maybe Packet
+    }
+ deriving Show
+
+data QueryMode = AskNot | AskAgain String | Ask
+ deriving (Show,Eq,Ord)
+
+getPassphrase :: GnuPGAgent -> QueryMode -> Query -> IO (Maybe String)
+getPassphrase agent ask (Query key uid masterkey) = do
+    let (er0,pr,desc) = prompts key uid masterkey
+        (er,askopt) = case ask of
+                        AskNot         -> (er0,"--no-ask")
+                        AskAgain ermsg -> (ermsg,"")
+                        Ask            -> (er0,"")
         cmd = "get_passphrase "++askopt++fingerprint key++" "++unwords (map percentPlusEscape [er,pr,desc])
-    putStrLn cmd
+    -- putStrLn cmd
     hPutStrLn (agentHandle agent) cmd
     r0 <- hGetLine (agentHandle agent)
+    -- putStrLn $ "agent says: " ++ r0
     case takeWhile (/=' ') r0 of
         "OK" -> hGetLine (agentHandle agent) >>= unhex . drop 3
             where
 #if defined(VERSION_memory)
                 unhex hx = case convertFromBase Base16 (S8.pack hx) of
-                            Left e -> return Nothing
+                            Left e -> do
+                                -- Useful for debugging but insecure generally ;)
+                                -- putStrLn $ "convertFromBase error for input "++show hx++": "++show e
+                                return Nothing
                             Right bs -> return $ Just $ S8.unpack bs
 #elif defined(VERSION_dataenc)
-                unhex hx = return $ fmap (map $ chr . fromIntegral) $ Base16.decode hx
+                unhex hx = maybe (return () {- putStrLn $ "dataenc error for input "++show hx -})
+                                 return
+                                 $ fmap (map $ chr . fromIntegral) $ Base16.decode hx
 #endif
         "ERR" -> return Nothing
 
diff --git a/lib/KeyRing.hs b/lib/KeyRing.hs
index ae2d14d..a055dad 100644
--- a/lib/KeyRing.hs
+++ b/lib/KeyRing.hs
@@ -47,7 +47,8 @@ module KeyRing
     , KeyFilter(..)
     -- * Results of a KeyRing Operation
     , KeyRingRuntime(..)
-    , MappedPacket(..)
+    , OriginMapped(..)
+    , MappedPacket
     , KeyDB
     , KeyData(..)
     , SubKey(..)
@@ -209,6 +210,7 @@ import Base58
 import FunctorToMaybe
 import DotLock
 import ProcessUtils (systemEnv, ExitCode(ExitFailure, ExitSuccess) )
+import GnuPGAgent as Agent
 
 -- DER-encoded elliptic curve ids
 -- nistp256_id  = 0x2a8648ce3d030107
@@ -429,6 +431,7 @@ data PassphraseSpec = PassphraseSpec
     }
     -- | Use this to carry pasphrases from a previous run.
     | PassphraseMemoizer PacketTranscoder
+    | PassphraseAgent
 
 instance Show PassphraseSpec where
     show (PassphraseSpec a b c) = "PassphraseSpec "++show (a,b,c)
@@ -1433,6 +1436,9 @@ doesInputFileExist ctx f = do
 -}
 
 
+-- | 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
@@ -1637,7 +1643,7 @@ buildKeyDB ctx grip0 keyring = do
         readw wk n = fmap (n,) (readPacketsFromWallet wk (ArgFile n))
 
     -- KeyRings (todo: KikiCondition reporting?)
-    (spilled,mwk,grip,accs,keys,unspilled) <- do
+    (spilled,mwk,grip,accs,keyqs,unspilled) <- do
 #if MIN_VERSION_containers(0,5,0)
         ringPackets <- Map.traverseWithKey readp ringMap
 #else
@@ -1660,32 +1666,69 @@ buildKeyDB ctx grip0 keyring = do
             -- | keys
             -- process ringPackets, and get a map of fingerprint info to 
             -- to a packet, remembering it's original file, access.
-            keys :: Map.Map KeyKey MappedPacket
+            keys :: Map.Map KeyKey (MappedPacket,Map.Map String [Packet])
             keys = Map.foldl slurpkeys Map.empty
                         $ Map.mapWithKey filterSecrets ringPackets
                     where
-                          filterSecrets f (_,Message ps) =
-                            filter (isSecretKey . packet)
-                                $ zipWith (mappedPacketWithHint fname) ps [1..]
-                            where fname = resolveForReport (Just ctx) f
-                          slurpkeys m ps = m `Map.union` Map.fromList ps'
-                            where ps' = zip (map (keykey . packet) ps) ps
+                          filterSecrets :: InputFile -> (a,Message) -> [[MappedPacket]]
+                          filterSecrets f (_,Message ps)  = keygroups
+                            -- filter (isSecretKey . packet) mps
+                            where
+                              mps = zipWith (mappedPacketWithHint fname) ps [1..]
+                              fname = resolveForReport (Just ctx) f
+                              keygroups = dropWhile (not . isSecretKey . packet . head)
+                                          $ groupBy (const $ not . isSecretKey . packet) mps
+                          slurpkeys :: (Map.Map KeyKey (MappedPacket,Map.Map String [Packet]))
+                                        -> [[MappedPacket]]
+                                        -> (Map.Map KeyKey (MappedPacket,Map.Map String [Packet]))
+                          slurpkeys m pss = Map.unionWith combineKeyKey m m2
+                            where
+                                m2 :: Map.Map KeyKey (MappedPacket, (Map.Map String [Packet]))
+                                m2 = Map.fromList $ map build pss
+                                    where
+                                        build ps = (kk,(kp,uidmap ps'))
+                                         where
+                                            (kpkt,ps') = splitAt 1 ps
+                                            kp = head kpkt
+                                            kk = keykey . packet $ kp
+                                combineKeyKey (mp,um) (mp2,um2) = (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
             -- | mwk 
             -- first master key matching the provided grip
             -- (the m is for "MappedPacket", wk for working key)
-            mwk :: Maybe MappedPacket
-            mwk = listToMaybe $ do
+            mwk = fst <$> mwkq
+
+            main_query = fromMaybe (Query MarkerPacket "anonymous1" Nothing) $ snd <$> mwkq
+
+            keyqs :: Map.Map KeyKey (OriginMapped Query)
+            keyqs = fmap (\(mp,us) -> mp { packet = main_query { queryPacket = packet mp} }) keys
+
+            mwkq :: Maybe (MappedPacket,Query)
+            mwkq = listToMaybe $ do
                     fp <- maybeToList grip
-                    let matchfp mp = not (is_subkey p) && matchpr fp p == fp
+                    let matchfp (mp,us)
+                         | not (is_subkey p) && matchpr fp p == fp = Just (mp,query p us)
+                         | otherwise                               = Nothing
                             where p = packet mp
-                    Map.elems $ Map.filter matchfp keys
+                                  -- TODO: check signature on UID packet?
+                                  -- TODO: custom queries for subkeys?
+                                  query p us = Query p
+                                                     (fromMaybe "" $ listToMaybe $ Map.keys us)
+                                                     Nothing -- No subkey queries for now.
+                    Map.elems $ Map.mapMaybe matchfp keys
+
             -- | accs
             -- file access(Sec | Pub) lookup table
             accs :: Map.Map InputFile Access
             accs = fmap (access . fst) ringPackets
-        return (spilled,mwk,grip,accs,keys,fmap snd unspilled)
+        return (spilled,mwk,grip,accs,keyqs,fmap snd unspilled)
 
-    transcode <- makeMemoizingDecrypter keyring ctx keys
+    putStrLn $ ppShow keyqs
+    transcode <- makeMemoizingDecrypter keyring ctx keyqs
     let doDecrypt = transcode (Unencrypted,S2K 100 "")
 
     let wk = fmap packet mwk
@@ -1707,6 +1750,9 @@ buildKeyDB ctx grip0 keyring = do
                 r <- performManipulations doDecrypt rt1 mwk manip
                 try r $ \(rt2,report) -> do
                 return $ KikiSuccess (report,rtKeyDB rt2)
+        -- XXX: Unspilled keys are not obtainable from rtKeyDB.
+        --  If the working key is marked non spillable, then how
+        --  would we look up it's UID and such?
 #if MIN_VERSION_containers(0,5,0)
         in fmap sequenceA $ Map.traverseWithKey trans spilled
 #else
@@ -2544,7 +2590,7 @@ writePEMKeys doDecrypt db exports = do
         return $ KikiSuccess (fname,stream,pun)
 
 makeMemoizingDecrypter :: KeyRingOperation -> InputFileContext
-                        -> Map.Map KeyKey MappedPacket
+                        -> Map.Map KeyKey (OriginMapped Query)
                         -> IO PacketTranscoder
 makeMemoizingDecrypter operation ctx keys = do
    if null chains then  do
@@ -2555,12 +2601,14 @@ makeMemoizingDecrypter operation ctx keys = do
     -- 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
@@ -2568,55 +2616,80 @@ makeMemoizingDecrypter operation ctx keys = 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))
-                                        $ opPassphrases operation
+                                        $ passspecs
         unkeysRef <- newIORef (Map.empty :: Map.Map KeyKey Packet)
-        return $ doDecrypt unkeysRef ({- pws `Map.union` -} pws2) defpw
+        return $ doDecrypt unkeysRef ({- pws `Map.union` -} pws2) defpw (not $ null agentspec)
     else let PassphraseMemoizer f = head chains
          in return f
  where
-    (chains,passspecs) = partition isChain $ opPassphrases operation
+    (chains,passspecs0) = partition isChain $ opPassphrases operation
                             where isChain (PassphraseMemoizer {}) = True
-                                  isChain _ = False
+                                  isChain _                       = False
+    (agentspec,passspecs) = partition isAgent $ opPassphrases operation
+                            where isAgent PassphraseAgent = True
+                                  isAgent _               = False
     doDecrypt :: IORef (Map.Map KeyKey Packet)
                 -> Map.Map FilePath (IO S.ByteString)
                 -> Maybe (IO S.ByteString)
+                -> Bool
                 -> (SymmetricAlgorithm,S2K)
                 -> MappedPacket
                 -> IO (KikiCondition Packet)
-    doDecrypt unkeysRef pws defpw (dest_alg,dest_s2k) mp0 = do
+    doDecrypt unkeysRef pws defpw agent_requested (dest_alg,dest_s2k) mp0 = do
         unkeys <- readIORef unkeysRef
-        let mp = fromMaybe mp0 $ do
-                    k <- Map.lookup kk keys
-                    return $ mergeKeyPacket "decrypt" mp0 k
+        let (mp,qry) = fromMaybe (mp0,Query (packet mp0) "anonymous2" Nothing) $ do
+                        k <- Map.lookup kk keys
+                        return (mergeKeyPacket "decrypt" mp0 (fmap queryPacket k), packet k)
             wk = packet mp0
             kk = keykey wk
             fs = Map.keys $ locations mp
 
             decryptIt [] = return BadPassphrase
-            decryptIt (getpw:getpws) = do
-                -- TODO: This function should use mergeKeyPacket to 
+            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.
-                pw <- getpw
+                --  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.
+                (pw,wants_retry) <- getpw (if count>1 then AskAgain "Bad pasphrase." else Ask,qry)
                 let wkun = fromMaybe wk $ do
                                 guard $ symmetric_algorithm (packet mp) /= Unencrypted
                                 decryptSecretKey pw (packet mp)
+
                 case symmetric_algorithm wkun of
+
                     Unencrypted -> do
                         writeIORef unkeysRef (Map.insert kk wkun unkeys)
                         ek <- if dest_alg==Unencrypted
                                 then return $ Just wkun
                                 else encryptSecretKey pw dest_s2k dest_alg wkun
                         case ek of
-                            Nothing   -> return $ BadPassphrase
-                            Just wken -> return $ KikiSuccess wken
-                    _                 -> decryptIt getpws
+                            Nothing | wants_retry && count<3 -> tries (count+1) getpw recurse
+                            Nothing                          -> recurse
+                            Just wken                        -> return $ KikiSuccess wken
+
+                    _ -> recurse
 
-            getpws = mapMaybe (`Map.lookup` pws) fs ++ maybeToList defpw
+            getpws = (map (const . fmap (,False)) $ mapMaybe (`Map.lookup` pws) fs ++ maybeToList defpw) ++ [ agentpw | agent_requested ]
+
+            agentpw (ask,qry) = do
+                    s <- session
+                    fromMaybe (return ("",False)) $ do
+                    s <- s
+                    Just $ do
+                    case ask of AskAgain _ -> clearPassphrase s (queryPacket qry)
+                                _          -> return ()
+                    mbpw <- getPassphrase s ask qry
+                    quit s
+                    return ( maybe "" S8.pack mbpw, True)
 
         if symmetric_algorithm wk == dest_alg
             && ( symmetric_algorithm wk == Unencrypted || s2k wk == dest_s2k )
@@ -3434,10 +3507,13 @@ data OriginFlags = OriginFlags
 
 type OriginMap = Map.Map FilePath OriginFlags
 
-data MappedPacket = MappedPacket
-        { packet :: Packet
+type MappedPacket = OriginMapped Packet
+data OriginMapped a = MappedPacket
+        { packet :: a
         , locations :: OriginMap
         } deriving Show
+instance Functor OriginMapped where
+    fmap f (MappedPacket x ls) = MappedPacket (f x) ls
 
 type TrustMap = Map.Map FilePath Packet
 type SigAndTrust = ( MappedPacket
@@ -3521,19 +3597,26 @@ onionName kd = (addr,name)
  where
     (addr,(name:_,_)) = getHostnames kd
 -}
+
+-- | 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++":"
-                              , fingerprint a
+                              , if isKey a then fingerprint a else ""
                               , PP.ppShow a
-                              , fingerprint b
+                              , if isKey b then fingerprint b else ""
                               , PP.ppShow b
                               ]
 
+-- | Merge two representations of the same key, prefering secret version
+-- because they have more information.
 mergeKeyPacket :: String -> MappedPacket -> MappedPacket -> MappedPacket
-mergeKeyPacket what key p = 
+mergeKeyPacket what key p =
     key { packet = minimumBy (keyCompare what) [packet key,packet p]
         , locations = Map.union (locations key) (locations p)
         }
diff --git a/lib/Kiki.hs b/lib/Kiki.hs
index a134680..25c98e2 100644
--- a/lib/Kiki.hs
+++ b/lib/Kiki.hs
@@ -41,6 +41,10 @@ import CommandLine
 import KeyRing
 import DotLock
 
+withAgent :: [PassphraseSpec] -> [PassphraseSpec]
+withAgent [] = [PassphraseAgent]
+withAgent ps = ps
+
 ciphername Unencrypted = "-"
 ciphername TripleDES   = "3des"
 ciphername (SymmetricAlgorithm w8) = "cipher-"++show w8
@@ -87,8 +91,8 @@ minimalOp cap = op
             [ ( HomeSec, streaminfo { access = Sec })
             , ( HomePub, streaminfo { access = Pub })
             ]
-        , opPassphrases = do pfile <- maybeToList (cap_passfd cap)
-                             return $ PassphraseSpec Nothing Nothing pfile
+        , opPassphrases = withAgent $ do pfile <- maybeToList (cap_passfd cap)
+                                         return $ PassphraseSpec Nothing Nothing pfile
         , opTransforms = []
         , opHome = cap_homespec cap
         }
@@ -155,6 +159,7 @@ importAndRefresh root cmn = do
             ctx = InputFileContext secring pubring
             passwordop = KeyRingOperation
                             { opFiles = Map.empty
+                            -- TODO: ask agent for new passphrase
                             , opPassphrases = do pfd <- maybeToList passfd
                                                  return $ PassphraseSpec Nothing Nothing pfd
                             , opHome = homespec
@@ -233,16 +238,16 @@ importAndRefresh root cmn = do
                 , ( ArgFile sshcpath, (peminfo 2048 "ssh-client") )
                 , ( ArgFile sshspath, (peminfo 2048 "ssh-server") )
                 ]
-             , opPassphrases = pwds ++ do pfd <- maybeToList passfd
-                                          return $ PassphraseSpec Nothing Nothing pfd
+             , opPassphrases = withAgent $ pwds ++ do pfd <- maybeToList passfd
+                                                      return $ PassphraseSpec Nothing Nothing pfd
              , opHome = homespec
              , opTransforms = []
              }
         -- doNothing = return ()
         nop = KeyRingOperation
             { opFiles = Map.empty
-            , opPassphrases = do pfd <- maybeToList passfd
-                                 return $ PassphraseSpec Nothing Nothing pfd
+            , opPassphrases = withAgent $ do pfd <- maybeToList passfd
+                                             return $ PassphraseSpec Nothing Nothing pfd
             , opHome=homespec, opTransforms = []
             }
     -- if bUnprivileged then doNothing else mkdirFor torpath