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 # mysql
  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 # postgresql
  28 # create table account (
  29 # id serial,
  30 # token char(128) unique not null,
  31 # cipher bytea not null,
  32 # primary key (id));
  33
  34 import base64
  35 from config import DBTYPE, DBHOST, DBPORT, DBUID, DBPWD, DBNAME
  36 from cryptonize import Cryptonize
  37
  38 class Singleton:
  39     def __init__(self, klass):
  40         self.klass = klass
  41         self.instance = None
  42
  43     def __call__(self, *args, **kwargs):
  44         if self.instance == None:
  45             self.instance = self.klass(*args, **kwargs)
  46         return self.instance
  47
  48 @Singleton
  49 class DB:
  50     conn = None
  51
  52     def get_connection(self):
  53         if self.conn is None:
  54             if DBTYPE == 'mysql':
  55                 #import MySQLdb
  56                 import MySQLdb
  57                 self.conn = MySQLdb.connect(host=DBHOST, port=DBPORT, user=DBUID, password=DBPWD, database=DBNAME)
  58             elif DBTYPE == 'postgresql':
  59                 import psycopg2
  60                 self.conn = psycopg2.connect(host=DBHOST, port=DBPORT, user=DBUID, password=DBPWD, dbname=DBNAME)
  61             else:
  62                 raise ValueError('{0}: Unsupported database'.format(DBTYPE))
  63         return self.conn
  64
  65     def __del__(self):
  66         if self.conn is not None:
  67             self.conn.close()
  68
  69 class DBInterface:
  70     @staticmethod
  71     def load_user(key):
  72         conn = DB().get_connection()
  73         cursor = conn.cursor()
  74         cursor.execute("select a.cipher from account a where token = '{0}'".format(key))
  75         row = cursor.fetchone()
  76         if row is None:
  77             obj = None
  78         else:
  79             c = Cryptonize()
  80             msg = base64.b64decode(row[0])
  81             obj = c.create_EncryptedMessage(msg)
  82         cursor.close()
  83
  84         return obj
  85
  86     @staticmethod
  87     def store_user(key, cipher):
  88         if DBTYPE == 'mysql':
  89             from MySQLdb import Error as DBError
  90         elif DBTYPE == 'postgresql':
  91             from psycopg2 import Error as DBError
  92         conn = DB().get_connection()
  93         cursor = conn.cursor()
  94         raw = base64.b64encode(cipher)
  95         try:
  96             cursor.execute("insert into account(token, cipher) values(%s, %s)", (key, raw))
  97             conn.commit()
  98         except DBError as e:
  99             print (e)
 100             conn.rollback()
 101             raise e
 102         finally:
 103             cursor.close()