some refactoring

This commit is contained in:
iou1name 2018-04-06 19:18:45 -04:00
parent d07aa07674
commit 0499adaddf

View File

@ -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()