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 # 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 # 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 # 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()