Refactor webservice to use SQLAlchemy and other good stuff

solvable/src/webservice/model.py

@@ -0,0 +1,22 @@
+from sqlalchemy import Column, Integer, String, create_engine
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import scoped_session, sessionmaker
+
+engine = create_engine('sqlite:///db.db', convert_unicode=True)
+db_session = scoped_session(sessionmaker(autocommit=False,
+                                         autoflush=False,
+                                         bind=engine))
+Base = declarative_base()
+Base.query = db_session.query_property()
+
+
+class User(Base):
+    __tablename__ = 'users'
+
+    id = Column(Integer, primary_key=True)
+    username = Column(String, nullable=False, unique=True)
+    passwordhash = Column(String, nullable=False)
+
+
+def init_db():
+    Base.metadata.create_all(bind=engine)

solvable/src/webservice/server.py

@@ -1,30 +1,12 @@
 from os import urandom, getenv
-from os.path import exists, join, dirname, realpath
 from hashlib import sha512
 
-import sqlite3
-from flask import Flask, render_template, request, g, session, url_for
-
-
-def setup_db(filename):
-    connection = sqlite3.connect(filename)
-    cur = connection.cursor()
-    cur.execute('''CREATE TABLE users
-                   (
-                      id INTEGER PRIMARY KEY,
-                      username TEXT NOT NULL,
-                      pwhash TEXT NOT NULL
-                   )
-                ''')
-    connection.commit()
-    connection.close()
+from flask import Flask, render_template, request, session, url_for
 
+from model import db_session, init_db, User
 
 BASEURL = getenv('BASEURL', '')
-DBFILE = join(dirname(realpath(__file__)), '.db.db')
-SALT = 'justsomerandombytes'.encode()
-if not exists(DBFILE):
-    setup_db(DBFILE)
+init_db()
 app = Flask(__name__)
 app.secret_key = urandom(32)
 
@@ -34,27 +16,23 @@ def get_url(endpoint):
 app.jinja_env.globals.update(get_url=get_url)
 
 
-@app.before_request
-def init_database():
-    g.db = sqlite3.connect(DBFILE)
-    g.db.row_factory = sqlite3.Row
+@app.teardown_appcontext
+def remove_db_session(exception=None):
+    db_session.remove()
 
 
-@app.teardown_request
-def close_database(exception):
-    db = getattr(g, 'db', None)
-    if db is not None:
-        db.close()
+def hash_password(password):
+    salt = 'justsomerandombytes'.encode()
+    return sha512(password.encode()+salt).hexdigest()
 
 
 @app.route('/', methods=['GET', 'POST'])
 def index():
     if request.method == 'POST':
-        cur = g.db.cursor()
-        cur.execute('SELECT * FROM users WHERE username=? AND pwhash=?', 
-                    [request.form['username'], sha512(request.form['password'].encode()+SALT).hexdigest()])
-        query = cur.fetchone()
-        if not query:
+        user = User.query.filter(User.username == request.form['username'], 
+                                 User.passwordhash == hash_password(request.form['password'])).first()
+
+        if not user:
             return render_template('login.html', alert='Invalid credentials!')
         else:
             session['logged_in'] = True
@@ -69,27 +47,26 @@ def index():
 @app.route('/register', methods=['GET', 'POST'])
 def register():
     if request.method == 'POST':
-        cur = g.db.cursor()
+        validate_register_fields(request)
 
-        if not request.form['username'] or not request.form['password'] or not request.form['passwordconfirm']:
-            return render_template('register.html', alert='You need to fill everything.')
-        if request.form['password'] != request.form['passwordconfirm']:
-            return render_template('register.html', alert='Passwords do not match! Please try again.')
-
-        cur.execute('SELECT * FROM users WHERE username=?', [request.form['username']])
-        if cur.fetchone() is not None:
+        if User.query.filter(User.username == request.form['username']).all():
             return render_template('register.html', alert='Username already in use.')
 
-        cur.execute('INSERT INTO users(username, pwhash) VALUES(?,?)', 
-                    [request.form['username'], 
-                     sha512(request.form['password'].encode()+SALT).hexdigest()])
-        g.db.commit()
+        db_session().add(User(username=request.form['username'],
+                              passwordhash=hash_password(request.form['password'])))
+        db_session().commit()
 
         return render_template('login.html', success='Account "{}" successfully registered. You can log in now!'.format(request.form['username']))
-
     return render_template('register.html')
 
 
+def validate_register_fields(request):
+    if not request.form['username'] or not request.form['password'] or not request.form['passwordconfirm']:
+        return render_template('register.html', alert='You need to fill everything.')
+    if request.form['password'] != request.form['passwordconfirm']:
+        return render_template('register.html', alert='Passwords do not match! Please try again.')
+
+
 @app.route('/logout')
 def logout():
     try:
@@ -117,4 +94,4 @@ def servererror(error):
 
 
 if __name__ == '__main__':
-    app.run(host='127.0.0.1', debug=False, port=6666)
+    app.run(host='127.0.0.1', debug=True, port=11111)

solvable/src/webservice/templates/internal.html

@@ -2,7 +2,7 @@
 
 {% block content %}
 <h2>Welcome, {{session.username}}</h2>
-<p>You have successfully logged in</p>
+<p>You have successfully logged in!</p>
 <a href="{{get_url('logout')}}" class="btn btn-default" >
     <span class="glyphicon glyphicon-log-out"></span> Logout
 </a>