working on changing from flask -> fastapi

1 files changed, 80 insertions, 59 deletions

diff --git a/app/routes.py b/app/routes.py
index a989c2e..fa73297 100644
--- a/app/routes.py
+++ b/app/routes.py
@@ -1,5 +1,8 @@
-import flask
+import fastapi
+from fastapi import Security, HTTPException
+from fastapi.security import APIKeyHeader
 import tabulate
+import pydantic
 import sqlalchemy
 
 from db import engine
@@ -12,47 +15,57 @@ from func.renew_link import renew_link
 from func.link_records import link_records
 from func.del_link_records import del_link_records
 
+class Newlink(pydantic.BaseModel):
+    redirect_link: str
 
-app = flask.Flask(__name__)
+app = fastapi.FastAPI()
+api_key_header = APIKeyHeader(name="X-API-Key")
 
-@app.route('/signup', methods=['GET'])
-def signup():
-    account_name = generate_account()
-    return flask.jsonify({'account_name': account_name})
+def check_api_key(api_key_header: str = Security(api_key_header)) -> str:
+    with engine.begin() as conn:
+        response = conn.execute(sqlalchemy.text("SELECT api_key FROM accounts WHERE api_key = :api_key"), {'api_key': api_key_header}).fetchone()
+    if response:
+        return response[0]
+    else:
+        raise HTTPException(
+            status_code=fastapi.status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or missing API key"
+        )
+
+
+@app.get("/signup")
+async def signup():
+    api_key = generate_account()
+    return {"api_key": api_key}
 
 
-@app.route('/newlink', methods=['POST'])
-@auth.login_required
-def newlink():
-    response = generate_link(flask.request, auth.current_user())
-    return flask.jsonify(msg=response[0]), response[1]
+@app.post("/newlink")
+async def newlink(newlink: Newlink, api_key: str = Security(check_api_key)):
+    data = generate_link(newlink.redirect_link, api_key)
+    if data:
+        return {"link": data[0], "expire_date": data[1]}
+    else:
+        raise HTTPException(
+            status_code=fastapi.status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail="Malformed redirect link provided"
+        )
 
 
-"""
-Return all links associated with an account
-"""
-@app.route('/links', methods=['POST'])
-@auth.login_required
-def links():
+@app.post("/links")
+async def links(api_key: str = Security(check_api_key)):
     with engine.begin() as conn:
-        links = conn.execute(sqlalchemy.text('SELECT link, expire_date FROM links WHERE owner = :owner'), [{'owner': auth.current_user()}]).fetchall()
+        links = conn.execute(sqlalchemy.text("SELECT link, expire_date FROM links WHERE owner = :owner"), [{"owner": api_key}]).fetchall()
 
-    string = ""
-    i = 1
+    response = []
     for link, expire_date in links:
-        string += f"{i}. {link} - Expires on {expire_date}\n"
-        i += 1
-    return string
+        response.append({"link": link, "expire_date": expire_date})
+    return response
 
 
-"""
-Return all records associated with an account, no matter the link
-"""
-@app.route('/records', methods=['POST'])
-@auth.login_required
-def records():
+@app.post("/records")
+async def records(api_key: str = Security(check_api_key)):
     with engine.begin() as conn:
-        records = conn.execute(sqlalchemy.text('SELECT timestamp, ip, location, browser, os, user_agent, isp FROM records WHERE owner = :owner'), [{'owner': auth.current_user()}]).fetchall()
+        records = conn.execute(sqlalchemy.text("SELECT timestamp, ip, location, browser, os, user_agent, isp FROM records WHERE owner = :owner"), [{"owner": api_key}]).fetchall()
 
     if not records:
         return flask.jsonify('No records found'), 200
@@ -60,40 +73,48 @@ def records():
     return tabulate.tabulate(records, headers=['Timestamp', 'IP', 'Location', 'Browser', 'OS', 'User Agent', 'ISP']), 200
 
 
-@app.route('/<link>', methods=['GET'])
-def link(link):
-    redirect_link = log(link, flask.request)
-    return flask.redirect(redirect_link)
+# """
+# Return all records associated with an account, no matter the link
+# """
+# @app.route('/records', methods=['POST'])
+# @auth.login_required
+# def records():
 
 
-@app.route('/<link>/delete', methods=['POST'])
-@auth.login_required
-def link_delete(link):
-    response = delete_link(link, auth.current_user())
-    return flask.jsonify(msg=response[0]), response[1]
+# @app.route('/<link>', methods=['GET'])
+# def link(link):
+#     redirect_link = log(link, flask.request)
+#     return flask.redirect(redirect_link)
 
 
-@app.route('/<link>/renew', methods=['POST'])
-@auth.login_required
-def renew_link(link):
-    response = renew_link(link, auth.current_user())
-    return flask.jsonify(msg=response[0]), response[1]
+# @app.route('/<link>/delete', methods=['POST'])
+# @auth.login_required
+# def link_delete(link):
+#     response = delete_link(link, auth.current_user())
+#     return flask.jsonify(msg=response[0]), response[1]
 
 
-@app.route('/<link>/records', methods=['POST'])
-@auth.login_required
-def records_link(link):
-    response = link_records(link, auth.current_user())
-    # If we jsonify the tabulate string it fucks it up, so we have to return
-    # it normally, this check does that
-    if response[0].startswith('Timestamp'):
-        return response[0], response[1]
-    else:
-        return flask.jsonify(msg=response[0]), response[1]
+# @app.route('/<link>/renew', methods=['POST'])
+# @auth.login_required
+# def renew_link(link):
+#     response = renew_link(link, auth.current_user())
+#     return flask.jsonify(msg=response[0]), response[1]
+
+
+# @app.route('/<link>/records', methods=['POST'])
+# @auth.login_required
+# def records_link(link):
+#     response = link_records(link, auth.current_user())
+#     # If we jsonify the tabulate string it fucks it up, so we have to return
+#     # it normally, this check does that
+#     if response[0].startswith('Timestamp'):
+#         return response[0], response[1]
+#     else:
+#         return flask.jsonify(msg=response[0]), response[1]
 
 
-@app.route('/<link>/delrecords', methods=['POST'])
-@auth.login_required
-def records_delete(link):
-    response = del_link_records(link, auth.current_user())
-    return flask.jsonify(msg=response[0]), response[1]
\ No newline at end of file
+# @app.route('/<link>/delrecords', methods=['POST'])
+# @auth.login_required
+# def records_delete(link):
+#     response = del_link_records(link, auth.current_user())
+#     return flask.jsonify(msg=response[0]), response[1]
\ No newline at end of file