multiple browser support, fixes #6

smorks · Mar 10, 2018 · 7f558f2 · 7f558f2
1 parent 94a92eb
commit 7f558f2
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 19 deletions.
diff --git a/KeePassNatMsg/Protocol/Action/Request.cs b/KeePassNatMsg/Protocol/Action/Request.cs
@@ -25,6 +25,8 @@ public static Request FromString(string s)
             return new Request((JObject)ReadFrom(rdr));
         }
 
+        public string ClientId => GetString("clientID");
+
         public string Action => GetString("action");
 
         public string Nonce => GetString("nonce");
@@ -41,7 +43,7 @@ public bool TryDecrypt()
         {
             try
             {
-                _msg = KeePassNatMsgExt.CryptoHelper.DecryptMessage(GetBytes("message"), GetBytes("nonce"));
+                _msg = KeePassNatMsgExt.CryptoHelper.DecryptMessage(ClientId, GetBytes("message"), GetBytes("nonce"));
                 return true;
             }
             catch (Exception)

diff --git a/KeePassNatMsg/Protocol/Action/Response.cs b/KeePassNatMsg/Protocol/Action/Response.cs
@@ -6,6 +6,7 @@ namespace KeePassNatMsg.Protocol.Action
     public class Response : JsonBase
     {
         private JsonBase _msg;
+        private string _clientId;
 
         public Response(Request req, bool createMessage)
         {
@@ -25,7 +26,7 @@ public string GetEncryptedResponse()
         {
             if (_msg != null)
             {
-                AddBytes("message", KeePassNatMsgExt.CryptoHelper.EncryptMessage(_msg.ToString(), Nonce));
+                AddBytes("message", KeePassNatMsgExt.CryptoHelper.EncryptMessage(_clientId, _msg.ToString(), Nonce));
             }
             return ToString();
         }
@@ -35,6 +36,7 @@ private void Init(Request req, bool createMessage)
             Add("action", new JValue(req.Action));
             AddBytes("nonce", Helper.GenerateNonce(req.NonceBytes));
             if (createMessage) CreateMessage();
+            _clientId = req.ClientId;
         }
 
         private void CreateMessage()

diff --git a/KeePassNatMsg/Protocol/Crypto/Helper.cs b/KeePassNatMsg/Protocol/Crypto/Helper.cs
@@ -1,38 +1,63 @@
 
 using KeePassNatMsg.Protocol.Action;
 using Newtonsoft.Json.Linq;
+using System.Collections.Generic;
 using System.Text;
 
 namespace KeePassNatMsg.Protocol.Crypto
 {
     public sealed class Helper
     {
-        private KeyPair _pair;
         private UTF8Encoding _utf8 = new UTF8Encoding(false);
+        private Dictionary<string, KeyPair> _clientKeys;
 
-        public byte[] ClientPublicKey { get; set; }
+        public Helper()
+        {
+            _clientKeys = new Dictionary<string, KeyPair>();
+        }
 
-        public JsonBase DecryptMessage(byte[] message, byte[] nonce)
+        public JsonBase DecryptMessage(string clientId, byte[] message, byte[] nonce)
         {
-            if (_pair == null) return null;
-            var data = TweetNaCl.CryptoBoxOpen(message, nonce, ClientPublicKey, _pair.PrivateKey);
-            return new JsonBase(JObject.Parse(_utf8.GetString(data)));
+            if (_clientKeys.ContainsKey(clientId))
+            {
+                var pair = _clientKeys[clientId];
+                var data = TweetNaCl.CryptoBoxOpen(message, nonce, pair.PublicKey, pair.PrivateKey);
+                return new JsonBase(JObject.Parse(_utf8.GetString(data)));
+            }
+            return null;
         }
 
-        public byte[] EncryptMessage(string msg, byte[] nonce)
+        public byte[] EncryptMessage(string clientId, string msg, byte[] nonce)
         {
-            return TweetNaCl.CryptoBox(_utf8.GetBytes(msg), nonce, ClientPublicKey, _pair.PrivateKey);
+            if (_clientKeys.ContainsKey(clientId))
+            {
+                var pair = _clientKeys[clientId];
+                return TweetNaCl.CryptoBox(_utf8.GetBytes(msg), nonce, pair.PublicKey, pair.PrivateKey);
+            }
+            return null;
         }
 
-        public KeyPair GenerateKeyPair()
+        public byte[] GenerateKeyPair(string clientId, byte[] clientPublicKey)
         {
-            _pair = new KeyPair();
-            return _pair;
+            var pair = new KeyPair();
+
+            if (_clientKeys.ContainsKey(clientId))
+            {
+                _clientKeys.Remove(clientId);
+            }
+
+            _clientKeys.Add(clientId, new KeyPair(pair.PrivateKey, clientPublicKey));
+
+            return pair.PublicKey;
         }
 
-        public void SetKeyPair(KeyPair pair)
+        public byte[] ClientPublicKey(string clientId)
         {
-            _pair = pair;
+            if (_clientKeys.ContainsKey(clientId))
+            {
+                return _clientKeys[clientId].PublicKey;
+            }
+            return null;
         }
 
         public static byte[] GenerateNonce(byte[] nonce) => TweetNaCl.Increment(nonce);

diff --git a/KeePassNatMsg/Protocol/Handlers.cs b/KeePassNatMsg/Protocol/Handlers.cs
@@ -111,7 +111,7 @@ private Response Associate(Request req)
             {
                 var msg = req.Message;
                 var keyBytes = msg.GetBytes("key");
-                if (keyBytes.SequenceEqual(KeePassNatMsgExt.CryptoHelper.ClientPublicKey))
+                if (keyBytes.SequenceEqual(KeePassNatMsgExt.CryptoHelper.ClientPublicKey(req.ClientId)))
                 {
                     var id = _ext.ShowConfirmAssociationDialog(msg.GetString("key"));
                     var resp = req.GetResponse();
@@ -135,10 +135,9 @@ private Response ChangePublicKeys(Request req)
             if (string.IsNullOrEmpty(publicKey))
                 return new ErrorResponse(req, ErrorType.ClientPublicKeyNotReceived);
 
-            crypto.ClientPublicKey = Convert.FromBase64String(publicKey);
-            var pair = crypto.GenerateKeyPair();
+            var serverPublicKey = crypto.GenerateKeyPair(req.ClientId, Convert.FromBase64String(publicKey));
             var resp = req.GetResponse(false);
-            resp.AddBytes("publicKey", pair.PublicKey);
+            resp.AddBytes("publicKey", serverPublicKey);
             resp.Add("version", KeePassNatMsgExt.GetVersion());
             resp.Add("success", "true");
             return resp;