From 6004ded914957910ab4bc8173db811098a9987a7 Mon Sep 17 00:00:00 2001
From: Michael Rasmussen <mir@datanom.net>
Date: Sun, 12 Aug 2018 00:31:54 +0200
Subject: [PATCH] Add feature to create tables in database

---
 .gitignore |  1 +
 db.py      | 96 ++++++++++++++++++++++++++++++++++++++++++++----------
 2 files changed, 80 insertions(+), 17 deletions(-)

diff --git a/.gitignore b/.gitignore
index 69fb7d4..303883f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,6 +13,7 @@ _ropeproject/
 *.bak
 *.rej
 *~
+*.db
 cur/
 tmp/
 __pycache__/
diff --git a/db.py b/db.py
index 9715b89..cabd8bf 100644
--- a/db.py
+++ b/db.py
@@ -18,25 +18,25 @@
 # along with SecureMail.  If not, see <https://www.gnu.org/licenses/>.
 
 # sqlite
-# create table account (
-# id int auto_increment,
-# token char(128) unique not null,
-# cipher text not null,
-# primary key (id));
-#
+sqlite_sql = """create table account (
+id int auto_increment,
+token char(128) unique not null,
+cipher text not null,
+primary key (id))"""
+
 # mysql
-# create table account (
-# id int auto_increment,
-# token char(128) unique not null,
-# cipher text not null,
-# primary key (id));
-#
+mysql_sql = """create table account (
+id int auto_increment,
+token char(128) unique not null,
+cipher text not null,
+primary key (id))"""
+
 # postgresql
-# create table account (
-# id serial,
-# token char(128) unique not null,
-# cipher bytea not null,
-# primary key (id));
+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
@@ -139,3 +139,65 @@ class DBInterface:
             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:
+            print (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()
-- 
2.39.2