some refactoring

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)
 
+	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
-
 	app.run(host='0.0.0.0', port=5000)
-
-else:
-	con, db = init()