82 lines
No EOL
2.8 KiB
Python
82 lines
No EOL
2.8 KiB
Python
import sqlite3
|
|
from os import path
|
|
import argparse
|
|
from flask import current_app, g
|
|
from passlib.hash import sha256_crypt
|
|
|
|
default_database_name = "sqlitedb"
|
|
|
|
def get_db():
|
|
current_app.config.get('DATABASE_PATH')
|
|
if 'db' not in g:
|
|
db_path = current_app.config.get('DATABASE_PATH')
|
|
if (db_path is None):
|
|
db_path = path.join(current_app.instance_path, current_app.config['DATABASE_NAME'])
|
|
g.db = sqlite3.connect(db_path, detect_types=sqlite3.PARSE_DECLTYPES)
|
|
g.db.row_factory = sqlite3.Row
|
|
|
|
return g.db
|
|
|
|
def close_db(e=None):
|
|
db = g.pop('db', None)
|
|
|
|
if db is not None:
|
|
db.close()
|
|
|
|
def init_db(db_path = None, schema_path = None):
|
|
if db_path is None:
|
|
db = get_db()
|
|
else:
|
|
db = sqlite3.connect(db_path, detect_types=sqlite3.PARSE_DECLTYPES)
|
|
|
|
if schema_path is None:
|
|
with current_app.open_resource('data/schema.sql') as f:
|
|
script = f.read().decode('UTF-8')
|
|
else:
|
|
with open(schema_path, "r") as f:
|
|
script = f.read()
|
|
|
|
db.executescript(script)
|
|
db.commit()
|
|
db.close()
|
|
|
|
def init_app(app):
|
|
app.teardown_appcontext(close_db)
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="DB Init ArgumentParser", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
parser.add_argument("-u", "--username", type=str, help="username of adming user", required=True)
|
|
parser.add_argument("-p", "--password", type=str, help="password of adming user", required=True)
|
|
parser.add_argument("-s", "--otp-secret", type=str, help="otp secret of admin user, python pyotp.random_base32()", required=False)
|
|
args = parser.parse_args()
|
|
config = vars(args)
|
|
username = config['username']
|
|
password = config['password']
|
|
otp_secret = config['otp_secret']
|
|
if (otp_secret is None):
|
|
import pyotp
|
|
otp_secret = pyotp.random_base32()
|
|
|
|
db_path = path.join('/server/instance/', default_database_name)
|
|
if path.exists(db_path):
|
|
print('Database already exists at {}. Will NOT override, if necessary first delete first then restart initialization!'.format(db_path))
|
|
exit(1)
|
|
schema_path = path.join(path.dirname(__file__), 'schema.sql')
|
|
init_db(db_path = db_path, schema_path = schema_path)
|
|
db = sqlite3.connect(db_path, detect_types=sqlite3.PARSE_DECLTYPES)
|
|
db.row_factory = sqlite3.Row
|
|
db_cursor = db.cursor()
|
|
|
|
hashed_password = sha256_crypt.hash(password)
|
|
sql = "INSERT INTO user(username, password, otp_secret, privileged, was_otp_verified)"\
|
|
"VALUES(:name, :pass, :otp, :privileged, :otp_verified)"
|
|
params = {
|
|
"name": username,
|
|
"pass": hashed_password,
|
|
"otp": otp_secret,
|
|
"privileged": True,
|
|
"otp_verified": False,
|
|
}
|
|
db_cursor.execute(sql, params)
|
|
db.commit()
|
|
db.close() |