Do not double print sql error

[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/>.

# 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, password=DBPWD, database=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()