X-Git-Url: http://git.datanom.net/securemail.git/blobdiff_plain/7ef0042650fe0a5e12373fd6ef438f138afe65b2..5bc1fc471eb0a72f4420b9189c990026e570d52a:/app/backend/db.py

diff --git a/app/backend/db.py b/app/backend/db.py
new file mode 100644
index 0000000..93eea09
--- /dev/null
+++ b/app/backend/db.py
@@ -0,0 +1,202 @@
+# -*- coding: utf-8 -*-
+ 
+# Copyright (c) 2018  Michael Rasmussen <mir@datanom.net>
+ 
+# This file is part of SecureMail.
+
+# SecureMail is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# SecureMail is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with SecureMail.  If not, see <https://www.gnu.org/licenses/>.
+
+# sqlite
+sqlite_sql = """create table account (
+id int auto_increment,
+token char(128) unique not null,
+cipher text not null,
+primary key (id))"""
+
+# mysql
+mysql_sql = """create table account (
+id int auto_increment,
+token char(128) unique not null,
+cipher text not null,
+primary key (id))"""
+
+# postgresql
+postgresql_sql = """create table account (
+id serial,
+token char(128) unique not null,
+cipher bytea not null,
+primary key (id))"""
+
+import base64
+from config import DBTYPE, DBNAME
+try:
+    from config import DBUID
+except ImportError:
+    DBUID = 'backend'
+try:
+    from config import DBPWD
+except ImportError:
+    DBPWD = 'clV77B2ZJQxr'
+try:
+    from config import DBHOST
+except ImportError:
+    DBHOST = 'localhost'
+try:
+    from config import DBPORT
+except ImportError:
+    if DBTYPE == 'mysql':
+        DBPORT = 3306
+    elif DBTYPE == 'postgresql':
+        DBPORT = 5432
+from cryptonize import Cryptonize
+
+class Singleton:
+    def __init__(self, klass):
+        self.klass = klass
+        self.instance = None
+
+    def __call__(self, *args, **kwargs):
+        if self.instance == None:
+            self.instance = self.klass(*args, **kwargs)
+        return self.instance
+
+@Singleton
+class DB:
+    conn = None
+
+    def get_connection(self):
+        if self.conn is None:
+            if DBTYPE == 'mysql':
+                import MySQLdb
+                self.conn = MySQLdb.connect(host=DBHOST, port=DBPORT, user=DBUID, passwd=DBPWD, db=DBNAME)
+            elif DBTYPE == 'postgresql':
+                import psycopg2
+                self.conn = psycopg2.connect(host=DBHOST, port=DBPORT, user=DBUID, password=DBPWD, dbname=DBNAME)
+            elif DBTYPE == 'sqlite':
+                import apsw
+                self.conn = apsw.Connection('./{0}.db'.format(DBNAME))
+            else:
+                raise ValueError('{0}: Unsupported database'.format(DBTYPE))
+        return self.conn
+        
+    def __del__(self):
+        if self.conn is not None:
+            self.conn.close()
+
+class DBInterface:
+    @staticmethod
+    def load_user(key):
+        conn = DB().get_connection()
+        cursor = conn.cursor()
+        cursor.execute("select a.cipher from account a where token = '{0}'".format(key))
+        row = cursor.fetchone()
+        if row is None:
+            obj = None
+        else:
+            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
+        elif DBTYPE == 'sqlite':
+            from apsw import Error as DBError
+        conn = DB().get_connection()
+        cursor = conn.cursor()
+        raw = base64.b64encode(cipher)
+        try:
+            if DBTYPE != 'sqlite':
+                cursor.execute("insert into account(token, cipher) values(%s, %s)", (key, raw))
+                conn.commit()
+            else:
+                cursor.execute('begin')
+                cursor.execute("insert into account(token, cipher) values(?, ?)", (key, raw))
+                cursor.execute('commit')
+        except DBError as e:
+            print (e)
+            if DBTYPE != 'sqlite':
+                conn.rollback()
+            else:
+                cursor.execute('rollback')
+            raise e
+        finally:
+            cursor.close()
+
+    @staticmethod
+    def create_database():
+        if DBTYPE == 'mysql':
+            from MySQLdb import Error as DBError
+        elif DBTYPE == 'postgresql':
+            from psycopg2 import Error as DBError
+        elif DBTYPE == 'sqlite':
+            from apsw import Error as DBError
+        conn = DB().get_connection()
+        cursor = conn.cursor()
+        try:
+            if DBTYPE != 'sqlite':
+                if DBTYPE == 'mysql':
+                    sql = mysql_sql
+                elif DBTYPE == 'postgresql':
+                    sql = postgresql_sql
+                cursor.execute(sql)
+                conn.commit()
+            else:
+                cursor.execute('begin')
+                cursor.execute(sqlite_sql)
+                cursor.execute('commit')
+        except DBError as e:
+            if DBTYPE != 'sqlite':
+                conn.rollback()
+            else:
+                cursor.execute('rollback')
+            raise e
+        finally:
+            cursor.close()
+
+def main():
+    from optparse import OptionParser
+    
+    usage = "usage: %prog [options] arg"
+    parser = OptionParser(usage)
+    parser.add_option("-c", "--create", action="store_true", dest="create",
+                      help="Create tables in database using config.py", default=False)
+    parser.add_option("-v", "--verbose", action="store_true", dest="verbose",
+                      help="Run in verbose mode", default=False)
+    (options, args) = parser.parse_args()
+    
+    if  options.create:
+        try:
+            if options.verbose:
+                print("Creating empty database")
+                print("Database Engine: {0}".format(DBTYPE))
+                if DBTYPE != 'sqlite':
+                    print("Database Host: {0}".format(DBHOST))
+                    print("Database Port: {0}".format(DBPORT))
+                else:
+                    print("Database File: ./{0}.db".format(DBNAME))
+            DBInterface.create_database()
+            print("Database created")
+        except Exception as e:
+            print("Creating database failed!")
+            print(e)
+        
+if __name__ == '__main__':
+    main()