2 files changed, 59 insertions, 8 deletions

diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 60019072..d23aca78 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -45,7 +45,7 @@ import Control.Concurrent.Lifted
 import GHC.Conc                  (labelThread)
 #endif
 
-import Crypto.Tox (zeros32,SecretKey,PublicKey)
+import Crypto.Tox (zeros32,SecretKey,PublicKey, generateSecretKey, toPublic, encodeSecret, decodeSecret)
 import Network.UPNP as UPNP
 import Network.Address hiding (NodeId, NodeInfo(..))
 import Network.Kademlia.Search
@@ -58,6 +58,7 @@ import Network.Kademlia.Routing as R
 import Data.Aeson as J (ToJSON, FromJSON)
 import qualified Data.Aeson as J
 import qualified Data.ByteString.Lazy as L
+import qualified Data.ByteString.Char8 as B
 import Control.Concurrent.Tasks
 import System.IO.Error
 import qualified Data.Serialize as S
@@ -283,7 +284,7 @@ data Session = Session
     , externalAddresses :: IO [SockAddr]
     , swarms            :: Mainline.SwarmsDatabase
     , toxkeys           :: TVar Tox.AnnouncedKeys
-    , keys              :: TVar [(SecretKey,PublicKey)]
+    , userkeys          :: TVar [(SecretKey,PublicKey)]
     , signalQuit        :: MVar ()
     }
 
@@ -295,6 +296,15 @@ clientSession s@Session{..} sock cnum h = do
         cmd0 action = action >> clientSession s sock cnum h
         switchNetwork dest = do hPutClient h ("Network: "++dest)
                                 clientSession s{netname=dest} sock cnum h
+        strp = B.unpack . fst . until snd dropEnd . (,False) . B.dropWhile isSpace . B.pack
+            where
+                dropEnd (x,_) =
+                    case B.unsnoc x of
+                                 Just (str,c) | isSpace c -> (str,False)
+                                 _ -> (x,True)
+    let mkrow :: (SecretKey, PublicKey) -> (String,String)
+        mkrow (a,b) | Just x <- encodeSecret a= (B.unpack x, show (Tox.key2id b))
+        mkrow _ = error (concat ["Assertion fail in 'mkrow' function at ", __FILE__, ":", show __LINE__])
     case (map toLower c,args) of
         ("stop", _) -> do hPutClient h "Terminating DHT Daemon."
                           hClose h
@@ -368,9 +378,44 @@ clientSession s@Session{..} sock cnum h = do
                                            let rs = [" ", show result]
                                            hPutClient h $ unlines rs
                           Left er -> hPutClient h er
-        ("k", "") -> cmd0 $ do
-            ks <- atomically $ readTVar keys
-            hPutClient h $ unlines $ map (mappend "  " . show . Tox.key2id . snd) ks
+        ("k", s) | ""           <- strp s -> cmd0 $ do
+                        ks <- atomically $ readTVar userkeys
+                        hPutClient h $ unlines $ map (mappend "  " . show . Tox.key2id . snd) ks
+                 | "gen"        <- strp s -> cmd0 $ do
+                        secret <- generateSecretKey
+                        let pubkey = toPublic secret
+                        oldks <- atomically $ do
+                            ks <- readTVar userkeys
+                            modifyTVar userkeys ((secret,pubkey):)
+                            return ks
+                        let asString = show . Tox.key2id
+                        hPutClient h $ unlines $ map (mappend "  " . show . Tox.key2id . snd) oldks
+                                                  ++ [mappend " *" . show . Tox.key2id $ pubkey]
+                 | "secrets"    <- strp s -> cmd0 $ do
+                        ks <- atomically $ readTVar userkeys
+                        hPutClient h . showReport $ map mkrow ks
+                 | ("add":secs) <- words s
+                 , mbSecs <- map (decodeSecret . B.pack) secs
+                 , all isJust mbSecs -> cmd0 $ do
+                        let f (Just b) = b
+                            f x = error (concat ["Assertion fail at ", __FILE__, ":", show __LINE__])
+                        let toPair x = (x,toPublic x)
+                            pairs = map (toPair . f) mbSecs
+                        oldks <- atomically $ readTVar userkeys
+                        atomically $ modifyTVar userkeys (pairs ++)
+                        hPutClient h $ unlines $ map (mappend "  " . show . Tox.key2id . snd) oldks
+                                                  ++ map (mappend " *" . show . Tox.key2id .snd) pairs
+                 | ("del":secs) <- words s
+                 , mbSecs <- map (decodeSecret . B.pack) secs
+                 , all isJust mbSecs -> cmd0 $ do
+                        let f (Just b) = b
+                            f x = error (concat ["Assertion fail at ", __FILE__, ":", show __LINE__])
+                        let toPair x = (x,toPublic x)
+                            pairs = map (toPair . f) mbSecs
+                        ks <- atomically $ do
+                            modifyTVar userkeys (filter (`notElem` pairs) )
+                            readTVar userkeys
+                        hPutClient h . showReport $ map mkrow ks
         ("g", s) | Just DHT{..} <- Map.lookup netname dhts
                     -> cmd0 $ do
                         -- arguments: method
@@ -628,12 +673,14 @@ main = do
 
     waitForSignal <- do
         signalQuit <- newEmptyMVar
+        userkeys0 <- atomically (newTVar [])
         let session = clientSession $ Session
                 { netname           = concat $ take 1 $ Map.keys dhts -- initial default DHT
                 , dhts              = dhts  -- all DHTs
                 , signalQuit        = signalQuit
                 , swarms            = swarms
                 , toxkeys           = keysdb
+                , userkeys          = userkeys0
                 , externalAddresses = liftM2 (++) btips toxips
                 }
         srv <- streamServer (withSession session) (SockAddrUnix "dht.sock")
diff --git a/src/Crypto/Tox.hs b/src/Crypto/Tox.hs
index b84e5df6..d6f63f18 100644
--- a/src/Crypto/Tox.hs
+++ b/src/Crypto/Tox.hs
@@ -13,6 +13,7 @@ module Crypto.Tox
     , getPublicKey
     , putPublicKey
     , SecretKey
+    , generateSecretKey
     , toPublic
     , SymmetricKey(..)
     , TransportCrypto(..)
@@ -359,7 +360,7 @@ getPublicKey = throwCryptoError . publicKey <$> S.getBytes 32
 putPublicKey :: PublicKey -> S.Put
 putPublicKey bs = S.putByteString $ BA.convert bs
 
-encodeSecret :: BA.ByteArrayAccess bin => bin -> Maybe C8.ByteString
+encodeSecret :: SecretKey -> Maybe C8.ByteString
 encodeSecret k = do
     (a,bs) <- BA.uncons (BA.convert k)
     (cs,c) <- unsnoc bs
@@ -369,7 +370,7 @@ encodeSecret k = do
         (ys,ds) = BA.splitAt 40 xs
     return $ BA.index ds 0 `BA.cons` ys `BA.snoc` BA.index ds 1
 
-decodeSecret :: C8.ByteString -> Maybe C8.ByteString
+decodeSecret :: C8.ByteString -> Maybe SecretKey
 decodeSecret k64 = do
     (ds0,ysds1) <- BA.uncons k64
     (ys,ds1) <- unsnoc ysds1
@@ -379,4 +380,7 @@ decodeSecret k64 = do
     (cs,a') <- unsnoc csa
     let a = shiftL (a' .&. 0x7c) 1
         c = shiftR c' 4 .|. (shiftL a' 4 .&. 0x30) .|. 0x40
-    return $ a `BA.cons` (cs `BA.snoc` c)
+    let r = a `BA.cons` (cs `BA.snoc` c)
+    case secretKey r of
+        CryptoPassed x -> Just x
+        _ -> Nothing