1 files changed, 19 insertions, 8 deletions

diff --git a/lib/KeyDB.hs b/lib/KeyDB.hs
index fd0a9ce..fc20b91 100644
--- a/lib/KeyDB.hs
+++ b/lib/KeyDB.hs
@@ -4,7 +4,7 @@ module KeyDB
     , SigAndTrust
     , SubKey(..)
     , KeyData(..)
-    , KeyDB
+    , KeyDB(..)
     , KeyGrip(..)
     , emptyKeyDB
     , keyData
@@ -68,6 +68,7 @@ data KeyData = KeyData
 
 
 newtype KeyGrip = KeyInt Int
+ deriving Eq
 
 fingerprintGrip :: Fingerprint -> KeyGrip
 fingerprintGrip (Fingerprint bs) =
@@ -109,10 +110,20 @@ kkData db = Map.toList (byKeyKey db)
 lookupKeyData :: KeyKey -> KeyDB -> Maybe KeyData
 lookupKeyData kk db = Map.lookup kk (byKeyKey db)
 
-lookupByGrip :: KeyGrip -> KeyDB -> [KeyData]
-lookupByGrip k db = mapMaybe (`Map.lookup` byKeyKey db)
-                    $ concat . maybeToList
-                    $ I.lookup k (byGrip db)
+lookupByGrip :: KeyGrip -> KeyDB -> [(MappedPacket,KeyData)]
+lookupByGrip k db = do
+    kk <- concat $ maybeToList $ I.lookup k (byGrip db)
+    case Map.lookup kk (byKeyKey db) of
+        Just kd | fingerprintGrip (fingerprint (packet $ keyMappedPacket kd)) == k -> [(keyMappedPacket kd, kd)]
+                | otherwise -> do
+                    sub <- associatedKeys kd
+                    guard (mpGrip sub == k)
+                    [ (sub, kd) ]
+        Nothing -> do
+            kd <- Map.elems (byKeyKey db)
+            sub <- associatedKeys kd
+            guard (mpGrip sub == k)
+            [ (sub, kd) ]
 
 transmute :: (Monad m, Monad kiki, Traversable kiki) =>
        ((KeyData, [info]) -> opcode -> m (kiki (KeyData, [info]))) -- ^ interpreter
@@ -137,14 +148,14 @@ associatedKeys kd = keyMappedPacket kd : [ k | SubKey k _ <- Map.elems (keySubKe
 alterKeyDB :: (Maybe KeyData -> Maybe KeyData) -> KeyKey -> KeyDB -> KeyDB
 alterKeyDB update kk db = db
     { byKeyKey = Map.alter update kk (byKeyKey db)
-    , byGrip = case Map.lookup kk (byKeyKey db) of
+    , byGrip = {- case Map.lookup kk (byKeyKey db) of
         Just _  -> byGrip db
-        Nothing -> case update Nothing of
+        Nothing -> -} case update Nothing of
             Just kd -> let go g m = I.alter (\case Nothing  -> Just [kk]
                                                    Just kks -> Just $ mergeL [kk] kks)
                                             g
                                             m
-                        in foldr go (byGrip db) $ map mpGrip $ associatedKeys kd
+                        in foldr go (byGrip db) $ map mpGrip $ filter (isKey . packet) $ associatedKeys kd
             Nothing -> byGrip db
     }