From 6a0146316be999e5bae9c37c69b9a63da2a8eb4b Mon Sep 17 00:00:00 2001
From: Timur Demin <me@tdem.in>
Date: Fri, 7 Jun 2019 23:45:23 +0500
Subject: [PATCH] Make errors a lot better

---
 project_amber/app.py          |  8 --------
 project_amber/errors.py       | 27 ++++++++++-----------------
 project_amber/helpers/auth.py | 25 ++++++++++++++++++-------
 3 files changed, 28 insertions(+), 32 deletions(-)

diff --git a/project_amber/app.py b/project_amber/app.py
index fd8a14f..a2d347a 100644
--- a/project_amber/app.py
+++ b/project_amber/app.py
@@ -1,7 +1,6 @@
 from json import dumps
 
 from flask import Flask
-from sqlalchemy.orm.exc import NoResultFound
 
 from project_amber.config import config
 from project_amber.db import db
@@ -27,10 +26,3 @@ def handle_HTTP_errors(e):
     return dumps({
         "message": e.message
     }), e.code
-
-# Hack.
-@app.errorhandler(NoResultFound)
-def handle_NoResultFound_errors(e):
-    return dumps({
-        "message": "Entity not found."
-    }), 404
diff --git a/project_amber/errors.py b/project_amber/errors.py
index 41c026d..f6043f8 100644
--- a/project_amber/errors.py
+++ b/project_amber/errors.py
@@ -2,8 +2,6 @@
 
 from project_amber.logging import logError
 
-# TODO: remake these into Werkzeug exceptions
-
 class HTTPError(Exception):
     """
     Base class for all possible errors.
@@ -24,42 +22,37 @@ class BadRequest(HTTPError):
     Exception class for payload data parsing errors.
     """
     code = HTTPStatus.BAD_REQUEST
-    message = "Bad request"
-    def __init__(self):
-        super().__init__(self.code, self.message)
+    def __init__(self, message="Bad request payload"):
+        super().__init__(self.code, message)
 
 class InternalServerError(HTTPError):
     """
     Exception class for DB errors. Probably going to be left unused.
     """
     code = HTTPStatus.INTERNAL_SERVER_ERROR
-    message = "Internal server error"
-    def __init__(self):
-        super().__init__(self.code, self.message)
+    def __init__(self, message="Internal error"):
+        super().__init__(self.code, message)
 
 class NotFound(HTTPError):
     """
     Exception class for entities not found.
     """
     code = HTTPStatus.NOT_FOUND
-    message = "Resource not found"
-    def __init__(self):
-        super().__init__(self.code, self.message)
+    def __init__(self, message="Entity not found"):
+        super().__init__(self.code, message)
 
 class NoAccess(HTTPError):
     """
     Exception class for restricted access areas.
     """
     code = HTTPStatus.FORBIDDEN
-    message = "Access denied"
-    def __init__(self):
-        super().__init__(self.code, self.message)
+    def __init__(self, message="Forbidden"):
+        super().__init__(self.code, message)
 
 class Unauthorized(HTTPError):
     """
     Exception class for login/auth check errors.
     """
     code = HTTPStatus.UNAUTHORIZED
-    message = "Unauthorized"
-    def __init__(self):
-        super().__init__(self.code, self.message)
+    def __init__(self, message="Unauthorized"):
+        super().__init__(self.code, message)
diff --git a/project_amber/helpers/auth.py b/project_amber/helpers/auth.py
index 3a385c7..fa6e8a8 100644
--- a/project_amber/helpers/auth.py
+++ b/project_amber/helpers/auth.py
@@ -6,7 +6,7 @@
 from flask import request
 
 from project_amber.db import db
-from project_amber.errors import Unauthorized, BadRequest
+from project_amber.errors import Unauthorized, BadRequest, NotFound, InternalServerError
 from project_amber.models.auth import User, Session
 
 class LoginUser:
@@ -34,9 +34,13 @@ def handleChecks() -> LoginUser:
         raise BadRequest
     token = request.headers.get("X-Auth-Token")
     if token is None:
-        raise Unauthorized
-    user_session = db.session.query(Session).filter_by(token=token).one()
-    user = db.session.query(User).filter_by(id=user_session.user).one()
+        raise Unauthorized("No X-Auth-Token present")
+    user_session = db.session.query(Session).filter_by(token=token).first()
+    if user_session is None:
+        raise Unauthorized("Invalid token")
+    user = db.session.query(User).filter_by(id=user_session.user).first()
+    if user is None:
+        raise InternalServerError("The user is missing")
     user_details = LoginUser(user.name, user.id, token)
     return user_details
 
@@ -55,7 +59,9 @@ def removeUser(uid: int) -> int:
     """
     Removes a user given their ID. Returns their ID on success.
     """
-    user = db.session.query(User).filter_by(id=uid).one()
+    user = db.session.query(User).filter_by(id=uid).first()
+    if user is None:
+        raise NotFound("User not found")
     db.session.delete(user)
     db.session.commit()
     return uid
@@ -73,7 +79,10 @@ def createSession(name: str, password: str) -> str:
     """
     Creates a new user session. Returns an auth token.
     """
-    user = db.session.query(User).filter_by(name=name).one()
+    user = db.session.query(User).filter_by(name=name).first()
+    if user is None:
+        raise Unauthorized # this may present no sense, but the app doesn't
+        # have to reveal the presence or absence of a user in the system
     if verifyPassword(user.id, password):
         token = sha256(gensalt()).hexdigest()
         session = Session(token=token, user=user.id, login_time=time())
@@ -86,7 +95,9 @@ def removeSession(token: str) -> str:
     """
     Removes a user session by token. Returns the token on success.
     """
-    session = db.session.query(Session).filter_by(token=token).one()
+    session = db.session.query(Session).filter_by(token=token).first()
+    if session is None:
+        raise NotFound
     db.session.delete(session)
     db.session.commit()
     return token