Basic framework finished

[securemail.git] / db.py

# -*- 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/>.

from config import DBTYPE, DBHOST, DBPORT, DBUID, DBPWD, DBNAME

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 mysql.connector
                self.conn = mysql.connector.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)
        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 id = '{0}'".format(key))
        row = cursor.fetchone()
        if row is None:
            obj = None
        else:
            obj = row[0].tobytes()
        cursor.close()
        
        return obj
        
    @staticmethod
    def store_user(key, cipher):
        conn = DB().get_connection()
        cursor = conn.cursor()
        cursor.execute("insert into account(id, cipher) values(%s, %s)", (key, cipher))
        conn.commit()
        cursor.close()