diff --git a/fileHost.py b/fileHost.py index 47ccb1f..187dd97 100755 --- a/fileHost.py +++ b/fileHost.py @@ -58,39 +58,36 @@ def init(): """ Initializes the application. """ - if not os.path.exists("secret_key"): - print("Error: secret_key file not found.") - print("If this is the first time the program is being run, run it \ - without your WSGI host first.") - os.exit() + os.makedirs(app.config.get("UPLOAD_DIR"), exist_ok=True) - with open("secret_key", "rb") as file: - secret_key = file.read() + if os.path.exists("secret_key"): + with open("secret_key", "rb") as file: + secret_key = file.read() + else: + secret_key = os.urandom(64) + with open("secret_key", "wb") as file: + file.write(secret_key) app.secret_key = secret_key - if not os.path.exists(app.config.get("DB_NAME")): - print("error: database not found.") - print("If this is the first time the program is being run, run it \ - without your WSGI host first.") - os.exit() - con = sqlite3.connect(app.config.get("DB_NAME")) db = con.cursor() + try: + db.execute("SELECT * FROM users").fetchone() + db.execute("SELECT * FROM uploads").fetchone() + except sqlite3.OperationalError: + db.execute("CREATE TABLE users(" + "id INTEGER PRIMARY KEY," + "username TEXT," + "pw_hash TEXT," + "admin BOOL DEFAULT FALSE)") + + db.execute("CREATE TABLE uploads(" + "filename TEXT, uploaded_by TEXT," + "uploaded_date INT DEFAULT (STRFTIME('%s', 'now')))") return con, db -def init_database(): - """ - Initializes appropriate tables in the database. - """ - db.execute("CREATE TABLE users(id INTEGER PRIMARY KEY, \ - username TEXT, pw_hash TEXT, admin BOOL DEFAULT FALSE)") - - db.execute("CREATE TABLE uploads(filename TEXT, uploaded_by TEXT, \ - uploaded_date INT DEFAULT (STRFTIME('%s', 'now')))") - - def add_user(username, password, admin="FALSE"): """ Adds a user to the database. @@ -292,29 +289,10 @@ def get_rand_chars(n): return "".join(chars) +con, db = init() +# TODO: make these not global variables if __name__ == "__main__": - os.makedirs(app.config.get("UPLOAD_DIR"), exist_ok=True) - initDB = not os.path.exists(app.config.get("DB_NAME")) - - con = sqlite3.connect(app.config.get("DB_NAME")) - db = con.cursor() - - if initDB: - print("Initializing new database and admin account.") - init_database() - import sys - add_user(sys.argv[1], sys.argv[2], "TRUE") - - if not os.path.exists("secret_key"): - secret_key = os.urandom(64) - with open("secret_key", "wb") as file: - file.write(secret_key) - else: - with open("secret_key", "rb") as file: - secret_key = file.read() - app.secret_key = secret_key + import sys + add_user(sys.argv[1], sys.argv[2], "TRUE") app.run(host='0.0.0.0', port=5000) - -else: - con, db = init()