db.py

   1 # -*- coding: utf-8 -*-
   2
   3 # Copyright (c) 2018  Michael Rasmussen <mir@datanom.net>
   4
   5 # This file is part of SecureMail.
   6
   7 # SecureMail is free software: you can redistribute it and/or modify
   8 # it under the terms of the GNU General Public License as published by
   9 # the Free Software Foundation, either version 3 of the License, or
  10 # (at your option) any later version.
  11 #
  12 # SecureMail is distributed in the hope that it will be useful,
  13 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15 # GNU General Public License for more details.
  16 #
  17 # You should have received a copy of the GNU General Public License
  18 # along with SecureMail.  If not, see <https://www.gnu.org/licenses/>.
  19
  20 # sqlite
  21 sqlite_sql = """create table account (
  22 id int auto_increment,
  23 token char(128) unique not null,
  24 cipher text not null,
  25 primary key (id))"""
  26
  27 # mysql
  28 mysql_sql = """create table account (
  29 id int auto_increment,
  30 token char(128) unique not null,
  31 cipher text not null,
  32 primary key (id))"""
  33
  34 # postgresql
  35 postgresql_sql = """create table account (
  36 id serial,
  37 token char(128) unique not null,
  38 cipher bytea not null,
  39 primary key (id))"""
  40
  41 import base64
  42 from config import DBTYPE, DBNAME
  43 try:
  44     from config import DBUID
  45 except ImportError:
  46     DBUID = 'backend'
  47 try:
  48     from config import DBPWD
  49 except ImportError:
  50     DBPWD = 'clV77B2ZJQxr'
  51 try:
  52     from config import DBHOST
  53 except ImportError:
  54     DBHOST = 'localhost'
  55 try:
  56     from config import DBPORT
  57 except ImportError:
  58     if DBTYPE == 'mysql':
  59         DBPORT = 3306
  60     elif DBTYPE == 'postgresql':
  61         DBPORT = 5432
  62 from cryptonize import Cryptonize
  63
  64 class Singleton:
  65     def __init__(self, klass):
  66         self.klass = klass
  67         self.instance = None
  68
  69     def __call__(self, *args, **kwargs):
  70         if self.instance == None:
  71             self.instance = self.klass(*args, **kwargs)
  72         return self.instance
  73
  74 @Singleton
  75 class DB:
  76     conn = None
  77
  78     def get_connection(self):
  79         if self.conn is None:
  80             if DBTYPE == 'mysql':
  81                 import MySQLdb
  82                 self.conn = MySQLdb.connect(host=DBHOST, port=DBPORT, user=DBUID, password=DBPWD, database=DBNAME)
  83             elif DBTYPE == 'postgresql':
  84                 import psycopg2
  85                 self.conn = psycopg2.connect(host=DBHOST, port=DBPORT, user=DBUID, password=DBPWD, dbname=DBNAME)
  86             elif DBTYPE == 'sqlite':
  87                 import apsw
  88                 self.conn = apsw.Connection('./{0}.db'.format(DBNAME))
  89             else:
  90                 raise ValueError('{0}: Unsupported database'.format(DBTYPE))
  91         return self.conn
  92
  93     def __del__(self):
  94         if self.conn is not None:
  95             self.conn.close()
  96
  97 class DBInterface:
  98     @staticmethod
  99     def load_user(key):
 100         conn = DB().get_connection()
 101         cursor = conn.cursor()
 102         cursor.execute("select a.cipher from account a where token = '{0}'".format(key))
 103         row = cursor.fetchone()
 104         if row is None:
 105             obj = None
 106         else:
 107             c = Cryptonize()
 108             msg = base64.b64decode(row[0])
 109             obj = c.create_EncryptedMessage(msg)
 110         cursor.close()
 111
 112         return obj
 113
 114     @staticmethod
 115     def store_user(key, cipher):
 116         if DBTYPE == 'mysql':
 117             from MySQLdb import Error as DBError
 118         elif DBTYPE == 'postgresql':
 119             from psycopg2 import Error as DBError
 120         elif DBTYPE == 'sqlite':
 121             from apsw import Error as DBError
 122         conn = DB().get_connection()
 123         cursor = conn.cursor()
 124         raw = base64.b64encode(cipher)
 125         try:
 126             if DBTYPE != 'sqlite':
 127                 cursor.execute("insert into account(token, cipher) values(%s, %s)", (key, raw))
 128                 conn.commit()
 129             else:
 130                 cursor.execute('begin')
 131                 cursor.execute("insert into account(token, cipher) values(?, ?)", (key, raw))
 132                 cursor.execute('commit')
 133         except DBError as e:
 134             print (e)
 135             if DBTYPE != 'sqlite':
 136                 conn.rollback()
 137             else:
 138                 cursor.execute('rollback')
 139             raise e
 140         finally:
 141             cursor.close()
 142
 143     @staticmethod
 144     def create_database():
 145         if DBTYPE == 'mysql':
 146             from MySQLdb import Error as DBError
 147         elif DBTYPE == 'postgresql':
 148             from psycopg2 import Error as DBError
 149         elif DBTYPE == 'sqlite':
 150             from apsw import Error as DBError
 151         conn = DB().get_connection()
 152         cursor = conn.cursor()
 153         try:
 154             if DBTYPE != 'sqlite':
 155                 if DBTYPE == 'mysql':
 156                     sql = mysql_sql
 157                 elif DBTYPE == 'postgresql':
 158                     sql = postgresql_sql
 159                 cursor.execute(sql)
 160                 conn.commit()
 161             else:
 162                 cursor.execute('begin')
 163                 cursor.execute(sqlite_sql)
 164                 cursor.execute('commit')
 165         except DBError as e:
 166             print (e)
 167             if DBTYPE != 'sqlite':
 168                 conn.rollback()
 169             else:
 170                 cursor.execute('rollback')
 171             raise e
 172         finally:
 173             cursor.close()
 174
 175 def main():
 176     from optparse import OptionParser
 177
 178     usage = "usage: %prog [options] arg"
 179     parser = OptionParser(usage)
 180     parser.add_option("-c", "--create", action="store_true", dest="create",
 181                       help="Create tables in database using config.py", default=False)
 182     parser.add_option("-v", "--verbose", action="store_true", dest="verbose",
 183                       help="Run in verbose mode", default=False)
 184     (options, args) = parser.parse_args()
 185
 186     if  options.create:
 187         try:
 188             if options.verbose:
 189                 print("Creating empty database")
 190                 print("Database Engine: {0}".format(DBTYPE))
 191                 if DBTYPE != 'sqlite':
 192                     print("Database Host: {0}".format(DBHOST))
 193                     print("Database Port: {0}".format(DBPORT))
 194                 else:
 195                     print("Database File: ./{0}.db".format(DBNAME))
 196             DBInterface.create_database()
 197             print("Database created")
 198         except Exception as e:
 199             print("Creating database failed!")
 200             print(e)
 201
 202 if __name__ == '__main__':
 203     main()