working on changing from flask -> fastapi

app/auth.py

@@ -10,7 +10,7 @@ auth = HTTPTokenAuth(scheme='Bearer')
 def verify_token(token):
     try:
         with engine.begin() as conn:
-            token = conn.execute(sqlalchemy.text('SELECT * FROM accounts WHERE account_name = :account_name'), [{'account_name': token}]).fetchone()
+            token = conn.execute(sqlalchemy.text('SELECT * FROM accounts WHERE api_key = :api_key'), [{'api_key': token}]).fetchone()
             return token[0]
     except TypeError:
         return False

app/db.py

@@ -13,7 +13,7 @@ def init_db():
         conn.execute(sqlalchemy.text(
             '''
             CREATE TABLE IF NOT EXISTS accounts (
-                account_name, PRIMARY KEY (account_name)
+                api_key, PRIMARY KEY (api_key)
             )
             '''
         ))
@@ -21,7 +21,7 @@ def init_db():
             '''
             CREATE TABLE IF NOT EXISTS links (
                 owner, link, redirect_link, expire_date,
-                FOREIGN KEY (owner) REFERENCES accounts(account_name), PRIMARY KEY (link)
+                FOREIGN KEY (owner) REFERENCES accounts(api_key), PRIMARY KEY (link)
             )
             '''
         ))

app/func/newlink.py

@@ -11,18 +11,9 @@ from db import engine
 Generate and return a new randomized link that is connected to the user
 Links are composed of 5 uppercase ASCII characters + numbers
 """
-def generate_link(request, owner):
-    content_type = request.headers.get('Content-Type')
-    if content_type == 'application/json':
-        try:
-            redirect_link = request.json['redirect_link']
-        except KeyError:
-            return 'Redirect link not provided', 400
-
-        if not validators.url(redirect_link):
-            return 'Redirect link is malformed. Please try again', 400
-    else:
-        return 'Content-Type not supported', 400
+def generate_link(redirect_link, owner):
+    if not validators.url(redirect_link):
+        return None
 
     with engine.begin() as conn:
         choices = string.ascii_uppercase + '1234567890'
@@ -36,4 +27,4 @@ def generate_link(request, owner):
             except exc.IntegrityError:
                 continue
 
-    return link, 200
+    return link, expire_date

app/func/signup.py

@@ -15,7 +15,7 @@ def generate_account():
         while True:
             try:
                 account_string = ''.join(random.choices(string.ascii_uppercase, k=20))
-                conn.execute(sqlalchemy.text('INSERT INTO accounts(account_name) VALUES(:account_name)'), [{'account_name': account_string}])
+                conn.execute(sqlalchemy.text('INSERT INTO accounts(api_key) VALUES(:api_key)'), [{'api_key': account_string}])
                 conn.commit()
                 break
             except exc.IntegrityError:

app/linklogger.py

@@ -1,19 +1,19 @@
-from routes import app
 from db import init_db
-import waitress
 import threading
 import schedule
 import time
+import uvicorn
 
 from func.remove_old_data import remove_old_data
 
+
 if __name__ == '__main__':
     init_db()
-    thread = threading.Thread(target=waitress.serve, args=(app,), kwargs={'port': 5252})
+    thread = threading.Thread(target=uvicorn.run("routes:app", host='127.0.0.1', port='5252'))
     thread.start()
     print('Server running on port 5252. Healthy.')
 
-    schedule.every().day.at('00:01').do(remove_old_data)
-    while True:
-        schedule.run_pending()
-        time.sleep(1)
+    # schedule.every().day.at('00:01').do(remove_old_data)
+    # while True:
+    #     schedule.run_pending()
+    #     time.sleep(1)

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})
-
-
-@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]
-
-
-"""
-Return all links associated with an account
-"""
-@app.route('/links', methods=['POST'])
-@auth.login_required
-def links():
+def check_api_key(api_key_header: str = Security(api_key_header)) -> str:
     with engine.begin() as conn:
-        links = conn.execute(sqlalchemy.text('SELECT link, expire_date FROM links WHERE owner = :owner'), [{'owner': auth.current_user()}]).fetchall()
+        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"
+        )
 
-    string = ""
-    i = 1
+
+@app.get("/signup")
+async def signup():
+    api_key = generate_account()
+    return {"api_key": api_key}
+
+
+@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"
+        )
+
+
+@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": api_key}]).fetchall()
+
+    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>/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]
+# @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]

requirements.txt

@@ -1,10 +1,11 @@
 Flask==3.0.0
 Flask-HTTPAuth==4.8.0
-waitress==2.1.2
 ip2location-io==1.0.0
 python-dotenv==1.0.0
 SQLAlchemy==2.0.27
 tabulate==0.9.0
 ua-parser==0.18.0
 validators==0.22.0
-schedule==1.2.1
+schedule==1.2.1
+
+# uvicorn, fastapi, pydantic