Tweak structure slightly (#2)

Dataflow breaks in the current setup, changing the code is quicker than a fix atm.

Author: ursachec

flask_webgoat/__init__.py

@@ -4,17 +4,26 @@
 
 from flask import Flask, g
 
+DB_FILENAME = 'database.db'
+
+
+def query_db(query, args = (), one=False, commit=False):
+ with sqlite3.connect(DB_FILENAME) as conn:
+ cur = conn.cursor().execute(query, args)
+ if commit:
+ conn.commit()
+ return cur.fetchone() if one else cur.fetchall()
+
 
 def create_app():
  app = Flask(__name__)
  app.secret_key = 'aeZ1iwoh2ree2mo0Eereireong4baitixaixu5Ee'
 
- db_filename = 'database.db'
- db_path = Path(db_filename)
+ db_path = Path(DB_FILENAME)
  if db_path.exists():
  db_path.unlink()
 
- conn = sqlite3.connect(db_filename)
+ conn = sqlite3.connect(DB_FILENAME)
  create_table_query = """CREATE TABLE IF NOT EXISTS user
  (id INTEGER PRIMARY KEY, username TEXT, password TEXT, access_level INTEGER)"""
  conn.execute(create_table_query)
@@ -25,13 +34,6 @@ def create_app():
  conn.commit()
  conn.close()
 
- def query_db(query, args = (), one=False, commit=False):
- with sqlite3.connect(db_filename) as conn:
- cur = conn.execute(query, args)
- if commit:
- conn.commit()
- return cur.fetchone() if one else cur.fetchall()
- app.query_db = query_db
 
  with app.app_context():
  from . import status

flask_webgoat/actions.py

@@ -28,12 +28,13 @@ def log_entry():
  return jsonify({'error': 'text parameter is required'})
 
  user_id = auth[0]
- user_dir = Path("data/" + str(user_id))
- if not user_dir.exists():
- user_dir.mkdir()
+ user_dir = "data/" + str(user_id)
+ user_dir_path = Path(user_dir)
+ if not user_dir_path.exists():
+ user_dir_path.mkdir()
 
  filename = filename_param + ".txt"
- path = user_dir / filename
+ path = Path(user_dir + "/" + filename)
  with path.open("w", encoding ="utf-8") as f:
  f.write(text_param)
  return jsonify({'success': True})

flask_webgoat/auth.py

@@ -3,6 +3,7 @@
 )
 from werkzeug.security import check_password_hash
 from flask import current_app as app
+from . import query_db
 
 bp = Blueprint('auth', __name__)
 
@@ -15,7 +16,7 @@ def login():
  return jsonify({'error': 'username and password parameter have to be provided'}), 400
 
  query = "SELECT id, username, access_level FROM user WHERE username = '%s' AND password = '%s'" % (username, password)
- result = app.query_db(query, [], True)
+ result = query_db(query, [], True)
  if result is None:
  return jsonify({'bad_login': True}), 400
  session['user_info'] = (result[0], result[1], result[2])

flask_webgoat/users.py

@@ -3,6 +3,7 @@
 )
 from werkzeug.security import check_password_hash
 from flask import current_app as app
+from . import query_db
 
 bp = Blueprint('users', __name__)
 
@@ -27,7 +28,7 @@ def create_user():
  query = "INSERT INTO user (username, password, access_level) VALUES ('%s', '%s', %d)" % (username, password, int(access_level))
 
  try:
- app.query_db(query, [], False, True)
+ query_db(query, [], False, True)
  return jsonify({'success': True})
  except sqlite3.Error as err:
  return jsonify({'error': 'could not create user:' + err})