X-Git-Url: http://git.datanom.net/securemail.git/blobdiff_plain/8c4f590c61472aa754a180e918ca5de7d1af5ad6..d65fab5a133cb7769beb24b7df81d356c12a6693:/db.py

diff --git a/db.py b/db.py
index e904fb1..1d545b1 100644
--- a/db.py
+++ b/db.py
@@ -17,7 +17,23 @@
 # You should have received a copy of the GNU General Public License
 # along with SecureMail.  If not, see <https://www.gnu.org/licenses/>.
 
+# mysql
+# create table account (
+# id int auto_increment,
+# token char(128) unique not null,
+# cipher blob not null,
+# primary key (id));
+#
+# postgresql
+# create table account (
+# id serial,
+# token char(128) unique not null,
+# cipher bytea not null,
+# primary key (id));
+
+import base64
 from config import DBTYPE, DBHOST, DBPORT, DBUID, DBPWD, DBNAME
+from cryptonize import Cryptonize
 
 class Singleton:
     def __init__(self, klass):
@@ -36,11 +52,14 @@ class DB:
     def get_connection(self):
         if self.conn is None:
             if DBTYPE == 'mysql':
-                import mysql.connector
-                self.conn = mysql.connector.connect(host=DBHOST, port=DBPORT, user=DBUID, password=DBPWD, database=DBNAME)
+                #import MySQLdb
+                import MySQLdb
+                self.conn = MySQLdb.connect(host=DBHOST, port=DBPORT, user=DBUID, password=DBPWD, database=DBNAME)
             elif DBTYPE == 'postgresql':
                 import psycopg2
                 self.conn = psycopg2.connect(host=DBHOST, port=DBPORT, user=DBUID, password=DBPWD, dbname=DBNAME)
+            else:
+                raise ValueError('{0}: Unsupported database'.format(DBTYPE))
         return self.conn
         
     def __del__(self):
@@ -52,20 +71,33 @@ class DBInterface:
     def load_user(key):
         conn = DB().get_connection()
         cursor = conn.cursor()
-        cursor.execute("select a.cipher from account a where id = '{0}'".format(key))
+        cursor.execute("select a.cipher from account a where token = '{0}'".format(key))
         row = cursor.fetchone()
         if row is None:
             obj = None
         else:
-            obj = row[0].tobytes()
+            c = Cryptonize()
+            msg = base64.b64decode(row[0])
+            obj = c.create_EncryptedMessage(msg)
         cursor.close()
         
         return obj
         
     @staticmethod
     def store_user(key, cipher):
+        if DBTYPE == 'mysql':
+            from MySQLdb import Error as DBError
+        elif DBTYPE == 'postgresql':
+            from psycopg2 import Error as DBError
         conn = DB().get_connection()
         cursor = conn.cursor()
-        cursor.execute("insert into account(id, cipher) values(%s, %s)", (key, cipher))
-        conn.commit()
-        cursor.close()
+        raw = base64.b64encode(cipher)
+        try:
+            cursor.execute("insert into account(token, cipher) values(%s, %s)", (key, raw))
+            conn.commit()
+        except DBError as e:
+            print (e)
+            conn.rollback()
+            raise e
+        finally:
+            cursor.close()