diff --git a/app/auth.py b/app/auth.py
index c9fe89b..aa278f2 100644
--- a/app/auth.py
+++ b/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
\ No newline at end of file
diff --git a/app/db.py b/app/db.py
index 454baae..c2ceb8b 100644
--- a/app/db.py
+++ b/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)
)
'''
))
diff --git a/app/func/newlink.py b/app/func/newlink.py
index ec45111..c726499 100644
--- a/app/func/newlink.py
+++ b/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
diff --git a/app/func/signup.py b/app/func/signup.py
index f7c19a2..275a14e 100644
--- a/app/func/signup.py
+++ b/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:
diff --git a/app/linklogger.py b/app/linklogger.py
index 640cff6..93ce72f 100644
--- a/app/linklogger.py
+++ b/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)
\ No newline at end of file
+ # schedule.every().day.at('00:01').do(remove_old_data)
+ # while True:
+ # schedule.run_pending()
+ # time.sleep(1)
\ No newline at end of file
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})
-
-
-@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('/', 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('//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('/', methods=['GET'])
+# def link(link):
+# redirect_link = log(link, flask.request)
+# return flask.redirect(redirect_link)
-@app.route('//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('//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('//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('//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('//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('//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('//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
diff --git a/requirements.txt b/requirements.txt
index 11fa0ed..bd4de63 100644
--- a/requirements.txt
+++ b/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
\ No newline at end of file
+schedule==1.2.1
+
+# uvicorn, fastapi, pydantic
\ No newline at end of file