Make backwards compatible with nacl 1.0.x

[securemail.git] / user.py

diff --git a/user.py b/user.py
index 3b4c051fb006078587a9adf111c4e63e15f20771..a4ab2f73e08fe3d5e7e20b95c17377be40ef9497 100644 (file)
--- a/user.py
+++ b/user.py
@@ -17,14 +17,12 @@
 # You should have received a copy of the GNU General Public License
 # along with SecureMail.  If not, see <https://www.gnu.org/licenses/>.
 
 # You should have received a copy of the GNU General Public License
 # along with SecureMail.  If not, see <https://www.gnu.org/licenses/>.
 

-try:
-    import cPickle as pickle
-except:
-    import pickle
+import pickle

 from db import DBInterface as DBI
 from cryptonize import Cryptonize
 from db import DBInterface as DBI
 from cryptonize import Cryptonize

+from nacl.public import PublicKey

 
 

-class NoSuchUserException(Exception):
+class NoSuchUser(Exception):

     pass
 
 class User:
     pass
 
 class User:


@@ -34,24 +32,48 @@ class User:
     def __init__(self, key=None):
         if key is not None:
             self.load(key)
     def __init__(self, key=None):
         if key is not None:
             self.load(key)

+        else:
+            self.pubkeys = {}

 
     def store(self, key):
         crypto = Cryptonize()
 
     def store(self, key):
         crypto = Cryptonize()

-        cipher = crypto.symmetric_encrypt(key, pickle.dumps(self.__dict__))
+        cipher = crypto.symmetric_encrypt(key, pickle.dumps(self))

         DBI.store_user(crypto.generate_hash(key), cipher)
 
     def load(self, key):
         crypto = Cryptonize()
         cipher = DBI.load_user(crypto.generate_hash(key))
         if cipher is None:
         DBI.store_user(crypto.generate_hash(key), cipher)
 
     def load(self, key):
         crypto = Cryptonize()
         cipher = DBI.load_user(crypto.generate_hash(key))
         if cipher is None:

-            raise NoSuchUserException('{0}: User not found'.format(key))
+            raise NoSuchUser('{0}: User not found'.format(key))

         plain = crypto.symmetric_decrypt(key, cipher)
         try:
             obj = pickle.loads(plain)
         plain = crypto.symmetric_decrypt(key, cipher)
         try:
             obj = pickle.loads(plain)

-            self.__dict__.update(obj)
+            self.__dict__.update(obj.__dict__)

         except pickle.UnpicklingError as e:
         except pickle.UnpicklingError as e:

-            raise NoSuchUserException(e)
-            
+            raise e
+    
+    def add_pubkey(self, email, key):
+        if email not in self.pubkeys:
+            self.pubkeys[email] = key.encode()
+        else:
+            raise KeyError('{0}: Exists'.format(email))
+    
+    def update_pubkey(self, email, key):
+        self.pubkeys[email] = key.encode()
+
+    def delete_pubkey(self, email):
+        if email in self.pubkeys:
+            del self.pubkeys[email]
+
+    def get_pubkey(self, email):
+        if email in self.pubkeys:
+            key = self.pubkeys[email]
+            key = PublicKey(key)
+        else:
+            key = None
+        
+        return key
+

     @property
     def name(self):
         return self._name
     @property
     def name(self):
         return self._name


@@ -62,25 +84,56 @@ class User:
 
     @property
     def email(self):
 
     @property
     def email(self):

-        return self.email
+        return self._email

 
     @email.setter
     def email(self, email):
         self._email = email
 
 
     @email.setter
     def email(self, email):
         self._email = email
 

+    @property
+    def pubkeys(self):
+        return self._pubkeys
+
+    @pubkeys.setter
+    def pubkeys(self, pubkeys):
+        if type(pubkeys) is not type({}):
+            raise ValueError('Not dictionary')
+        self._pubkeys = pubkeys

 
 if __name__ == '__main__':
     try:
         u = User('test')
         for attr, value in u.__dict__.items():
             print ('{0}: {1}'.format(attr, value))
 
 if __name__ == '__main__':
     try:
         u = User('test')
         for attr, value in u.__dict__.items():
             print ('{0}: {1}'.format(attr, value))

+        print ('{0} - {1} - {2}'.format(u.name, u.email, u.pubkeys))
+        key = ''
+        for i in range(40):
+            key += '{0}'.format(i)
+        u = User()
+        u.name = 'testname1'
+        u.email = 'testname1@securemail.icu'
+        u.pubkeys = {'test': 'some test', 'test1': 'some test 1'}
+        try:
+            u.store(key)
+        except:
+            u = User(key)
+            for attr, value in u.__dict__.items():
+                print ('{0}: {1}'.format(attr, value))
+            print ('{0} - {1} - {2}'.format(u.name, u.email, u.pubkeys))

         c = Cryptonize()
         c = Cryptonize()

-        key = 'æselØre' #c.get_random_key()
-        cipher = c.symmetric_encrypt(key, pickle.dumps(u))
-        obj = pickle.loads(c.symmetric_decrypt(key, cipher))
-        for attr, value in obj.__dict__.items():
-            print ('{0}: {1}'.format(attr, value))
-    except NoSuchUserException:
+        keypair1 = c.get_key_pair()
+        keypair2 = c.get_key_pair()
+        try:
+            u.add_pubkey('test', keypair2[1])
+        except KeyError:
+            u.update_pubkey('test', keypair2[1])
+        message = "Kill all humans æøåÅØÆ"
+        print ("Message to encrypt: {0}".format(message))
+        encrypted = c.asymmetric_encrypt(keypair1[0], u.get_pubkey('test'), message)
+        print ("Message encrypted: {0}".format(encrypted))
+        plaintext = c.asymmetric_decrypt(keypair2[0], keypair1[1], encrypted)
+        print("Message decrypted: {0}".format(plaintext.decode()))
+    except NoSuchUser:

         u = User()
         u.name = 'testname'
         u.email = 'testname@securemail.icu'
         u = User()
         u.name = 'testname'
         u.email = 'testname@securemail.icu'