14 files changed, 223 insertions, 156 deletions

diff --git a/app/auth.py b/app/auth.py
deleted file mode 100644
index aa278f2..0000000
--- a/app/auth.py
+++ /dev/null
@@ -1,16 +0,0 @@
-from flask_httpauth import HTTPTokenAuth
-import sqlalchemy
-
-from db import engine
-
-
-auth = HTTPTokenAuth(scheme='Bearer')
-
-@auth.verify_token
-def verify_token(token):
-    try:
-        with engine.begin() as conn:
-            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/check_api_key.py b/app/check_api_key.py
new file mode 100644
index 0000000..d8b92f4
--- /dev/null
+++ b/app/check_api_key.py
@@ -0,0 +1,22 @@
+import fastapi
+from fastapi import Security, HTTPException
+from fastapi.security import APIKeyHeader
+import sqlalchemy
+
+from db import engine
+
+"""
+Make sure the provided API key is valid
+"""
+api_key_header = APIKeyHeader(name="X-API-Key")
+
+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 keys 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"
+        )
diff --git a/app/db.py b/app/db.py
index c2ceb8b..1750472 100644
--- a/app/db.py
+++ b/app/db.py
@@ -12,7 +12,7 @@ def init_db():
     with engine.begin() as conn:
         conn.execute(sqlalchemy.text(
             '''
-            CREATE TABLE IF NOT EXISTS accounts (
+            CREATE TABLE IF NOT EXISTS keys (
                 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(api_key), PRIMARY KEY (link)
+                FOREIGN KEY (owner) REFERENCES keys(api_key), PRIMARY KEY (link)
             )
             '''
         ))
diff --git a/app/func/generate_api_key.py b/app/func/generate_api_key.py
new file mode 100644
index 0000000..a40c96a
--- /dev/null
+++ b/app/func/generate_api_key.py
@@ -0,0 +1,23 @@
+import sqlalchemy
+from sqlalchemy import exc
+import random
+import string
+
+from db import engine
+
+"""
+Generate and return a randomized API key string for the user
+Keys are composed of 20 uppercase ASCII characters
+"""
+def generate_api_key():
+    with engine.begin() as conn:
+        while True:
+            try:
+                api_key_string = ''.join(random.choices(string.ascii_uppercase, k=20))
+                conn.execute(sqlalchemy.text('INSERT INTO keys(api_key) VALUES(:api_key)'), [{'api_key': api_key_string}])
+                conn.commit()
+                break
+            except exc.IntegrityError:
+                continue
+
+    return api_key_string
\ No newline at end of file
diff --git a/app/func/delete_link.py b/app/func/link/delete.py
index c036af3..97b696e 100644
--- a/app/func/delete_link.py
+++ b/app/func/link/delete.py
@@ -10,11 +10,11 @@ def delete_link(link, owner):
         try:
             link_owner = conn.execute(sqlalchemy.text('SELECT owner FROM links WHERE link = :link'), [{'link': link}]).fetchone()[0]
         except TypeError:
-            return 'Link does not exist', 200
+            return 404
 
     if owner == link_owner:
         with engine.begin() as conn:
             conn.execute(sqlalchemy.text('DELETE FROM links WHERE link = :link'), [{'link': link}])
-            return 'Link has been deleted', 200
+            return link
     else:
-        return 'You are not the owner of this link', 401
\ No newline at end of file
+        return 401
\ No newline at end of file
diff --git a/app/func/del_link_records.py b/app/func/link/delrecords.py
index fbbce0e..d82bfa5 100644
--- a/app/func/del_link_records.py
+++ b/app/func/link/delrecords.py
@@ -5,16 +5,16 @@ from db import engine
 """
 Delete all of the IP log records that are associated with a specific link
 """
-def del_link_records(link, owner):
+def delete_link_records(link, owner):
     with engine.begin() as conn:
         try:
             link_owner = conn.execute(sqlalchemy.text('SELECT owner FROM links WHERE link = :link'), [{'link': link}]).fetchone()[0]
         except TypeError:
-            return 'Link does not exist', 200
+            return 404
 
     if owner == link_owner:
         with engine.begin() as conn:
             conn.execute(sqlalchemy.text('DELETE FROM records WHERE link = :link'), [{'link': link}])
-            return 'Link records have been deleted', 200
+            return link
     else:
-        return 'You are not the owner of this link', 401
\ No newline at end of file
+        return 401
\ No newline at end of file
diff --git a/app/func/link_records.py b/app/func/link/records.py
index a29f8dd..56bb6d2 100644
--- a/app/func/link_records.py
+++ b/app/func/link/records.py
@@ -1,24 +1,23 @@
 import sqlalchemy
-import tabulate
 
 from db import engine
 
 """
 Retrieve all records associated with a specific link
 """
-def link_records(link, owner):
+def get_link_records(link, owner):
     with engine.begin() as conn:
         try:
             link_owner = conn.execute(sqlalchemy.text('SELECT owner FROM links WHERE link = :link'), [{'link': link}]).fetchone()[0]
         except TypeError:
-            return 'Link does not exist', 200
+            return 404
 
     if owner == link_owner:
         with engine.begin() as conn:
             records = conn.execute(sqlalchemy.text('SELECT timestamp, ip, location, browser, os, user_agent, isp FROM records WHERE owner = :owner and link = :link'), [{'owner': owner, 'link': link}]).fetchall()
         if not records:
-            return 'No records are associated with this link', 200
+            return 204
     else:
-        return 'You are not the owner of this link', 401
+        return 401
 
-    return tabulate.tabulate(records, headers=['Timestamp', 'IP', 'Location', 'Browser', 'OS', 'User Agent', 'ISP']), 200
\ No newline at end of file
+    return records
\ No newline at end of file
diff --git a/app/func/renew_link.py b/app/func/link/renew.py
index 9d31c33..bcf0550 100644
--- a/app/func/renew_link.py
+++ b/app/func/link/renew.py
@@ -12,12 +12,12 @@ def renew_link(link, owner):
         try:
             link_owner = conn.execute(sqlalchemy.text('SELECT owner FROM links WHERE link = :link'), [{'link': link}]).fetchone()[0]
         except TypeError:
-            return 'Link does not exist', 200
+            return 404
 
     if owner == link_owner:
         with engine.begin() as conn:
             expire_date = datetime.datetime.date(datetime.datetime.now()) + datetime.timedelta(days=7)
             conn.execute(sqlalchemy.text('UPDATE links SET expire_date = :expire_date WHERE link = :link'), [{'expire_date': expire_date, 'link': link}])
-            return f'Link renewed, now expires on {expire_date}', 200
+            return link, expire_date
     else:
-        return 'You are not the owner of this link', 401
\ No newline at end of file
+        return 401
\ No newline at end of file
diff --git a/app/func/log.py b/app/func/log.py
index 1199f43..da6594a 100644
--- a/app/func/log.py
+++ b/app/func/log.py
@@ -39,7 +39,7 @@ ipgeolocation = ip2locationio.IPGeolocation(configuration)
 """
 Create a new log record whenever a link is visited
 """
-def log(link, request):
+def log(link, ip, user_agent):
     with engine.begin() as conn:
         try:
             redirect_link, owner = conn.execute(sqlalchemy.text('SELECT redirect_link, owner FROM links WHERE link = :link'), [{'link': link}]).fetchone()
@@ -50,7 +50,7 @@ def log(link, request):
         if ip_to_location == 'TRUE':
             # Get IP to GEO via IP2Location.io
             try:
-                data = ipgeolocation.lookup(request.access_route[-1])
+                data = ipgeolocation.lookup(ip)
                 location = f'{data["country_name"]}, {data["city_name"]}'
                 isp = data['as']
             # Fatal error, API key is invalid or out of requests, quit
@@ -63,8 +63,6 @@ def log(link, request):
             isp = '-'
 
         timestamp = datetime.datetime.now()
-        ip = request.access_route[-1]
-        user_agent = request.user_agent.string
         ua_string = user_agent_parser.Parse(user_agent)
         browser = ua_string['user_agent']['family']
         os = f'{ua_string["os"]["family"]} {ua_string["os"]["major"]}'
diff --git a/app/func/newlink.py b/app/func/newlink.py
index c726499..06776e6 100644
--- a/app/func/newlink.py
+++ b/app/func/newlink.py
@@ -13,7 +13,7 @@ Links are composed of 5 uppercase ASCII characters + numbers
 """
 def generate_link(redirect_link, owner):
     if not validators.url(redirect_link):
-        return None
+        return 422
 
     with engine.begin() as conn:
         choices = string.ascii_uppercase + '1234567890'
diff --git a/app/func/remove_old_data.py b/app/func/remove_old_data.py
index 7380fc5..96d08fa 100644
--- a/app/func/remove_old_data.py
+++ b/app/func/remove_old_data.py
@@ -17,7 +17,8 @@ def remove_old_data():
         link = row.link
         delete_links.append({'link': link})
 
-    with engine.begin() as conn:
-        conn.execute(sqlalchemy.text('DELETE FROM links WHERE link = :link'), delete_links)
-        conn.execute(sqlalchemy.text('DELETE FROM records WHERE link = :link'), delete_links)
-        conn.commit()
+    if delete_links:
+        with engine.begin() as conn:
+            conn.execute(sqlalchemy.text('DELETE FROM links WHERE link = :link'), delete_links)
+            conn.execute(sqlalchemy.text('DELETE FROM records WHERE link = :link'), delete_links)
+            conn.commit()
diff --git a/app/func/signup.py b/app/func/signup.py
deleted file mode 100644
index 275a14e..0000000
--- a/app/func/signup.py
+++ /dev/null
@@ -1,24 +0,0 @@
-import sqlalchemy
-from sqlalchemy import exc
-import random
-import string
-
-from db import engine
-
-"""
-Generate and return a randomized account string for the user
-Account strings function as API authenticaton keys and are composed
-of 20 uppercase ASCII characters
-"""
-def generate_account():
-    with engine.begin() as conn:
-        while True:
-            try:
-                account_string = ''.join(random.choices(string.ascii_uppercase, k=20))
-                conn.execute(sqlalchemy.text('INSERT INTO accounts(api_key) VALUES(:api_key)'), [{'api_key': account_string}])
-                conn.commit()
-                break
-            except exc.IntegrityError:
-                continue
-
-    return account_string
\ No newline at end of file
diff --git a/app/linklogger.py b/app/linklogger.py
index 93ce72f..388ede7 100644
--- a/app/linklogger.py
+++ b/app/linklogger.py
@@ -1,19 +1,9 @@
 from db import init_db
-import threading
-import schedule
-import time
 import uvicorn
 
-from func.remove_old_data import remove_old_data
+from routes import app
 
 
 if __name__ == '__main__':
     init_db()
-    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
+    server = uvicorn.run(app=app, host="0.0.0.0", port=5252)
\ No newline at end of file
diff --git a/app/routes.py b/app/routes.py
index fa73297..da8a502 100644
--- a/app/routes.py
+++ b/app/routes.py
@@ -1,120 +1,194 @@
 import fastapi
-from fastapi import Security, HTTPException
-from fastapi.security import APIKeyHeader
-import tabulate
+from fastapi import Security, HTTPException, Request
 import pydantic
 import sqlalchemy
 
 from db import engine
-from auth import auth
-from func.signup import generate_account
+from check_api_key import check_api_key
+from func.generate_api_key import generate_api_key
 from func.newlink import generate_link
 from func.log import log
-from func.delete_link import delete_link
-from func.renew_link import renew_link
-from func.link_records import link_records
-from func.del_link_records import del_link_records
+from func.link.delete import delete_link
+from func.link.renew import renew_link
+from func.link.records import get_link_records
+from func.link.delrecords import delete_link_records
+from func.remove_old_data import remove_old_data
+
+from apscheduler.schedulers.background import BackgroundScheduler
+from contextlib import asynccontextmanager
+
+
+@asynccontextmanager
+async def lifespan(app: fastapi.FastAPI):
+    # Create the scheduler
+    scheduler = BackgroundScheduler()
+    scheduler.add_job(remove_old_data, "cron", hour="0", minute="01")
+    scheduler.start()
+    yield
 
 class Newlink(pydantic.BaseModel):
     redirect_link: str
 
-app = fastapi.FastAPI()
-api_key_header = APIKeyHeader(name="X-API-Key")
+app = fastapi.FastAPI(lifespan=lifespan)
 
-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()
+@app.post("/api/getapikey")
+async def get_api_key():
+    """
+    Create a new API key
+    """
+    api_key = generate_api_key()
     return {"api_key": api_key}
 
 
-@app.post("/newlink")
+@app.post("/api/genlink")
 async def newlink(newlink: Newlink, api_key: str = Security(check_api_key)):
+    """
+    Generate a new link that will redirect to the specified URL and log IPs in the middle
+    """
     data = generate_link(newlink.redirect_link, api_key)
-    if data:
-        return {"link": data[0], "expire_date": data[1]}
-    else:
+    if data == 422:
         raise HTTPException(
             status_code=fastapi.status.HTTP_422_UNPROCESSABLE_ENTITY,
             detail="Malformed redirect link provided"
         )
 
+    return {"link": data[0], "expire_date": data[1]}
 
-@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:
-        response.append({"link": link, "expire_date": expire_date})
-    return response
-
-
-@app.post("/records")
+"""
+Return all records associated with an API key, no matter the link
+"""
+@app.get("/api/records")
 async def records(api_key: str = Security(check_api_key)):
+    """
+    Get ALL IP logs records for every link tied to your 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": api_key}]).fetchall()
 
     if not records:
-        return flask.jsonify('No records found'), 200
+        return {"records": "No records are associated with this API key"}
+
+    response = []
+    for timestamp, ip, location, browser, os, user_agent, isp in records:
+        response.append({"timestamp": timestamp, "ip": ip, "location": location, "browser": browser, "os": os, "user_agent": user_agent, "isp": isp})
+
+    return response
 
-    return tabulate.tabulate(records, headers=['Timestamp', 'IP', 'Location', 'Browser', 'OS', 'User Agent', 'ISP']), 200
 
+@app.get("/{link}")
+def link(link, request: Request):
+    ip = request.client.host
+    user_agent = request.headers.get("user-agent")
+    redirect_link = log(link, ip, user_agent)
+    return fastapi.responses.RedirectResponse(url=redirect_link)
 
-# """
-# Return all records associated with an account, no matter the link
-# """
-# @app.route('/records', methods=['POST'])
-# @auth.login_required
-# def records():
 
+"""
+Return all links associated with an API key
+"""
+@app.get("/api/links")
+async def links(api_key: str = Security(check_api_key)):
+    """
+    Retrieve all links that are currently tied to your 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()
+
+    if not links:
+        return {"links": "No links are associated with this API key"}
+
+    response = []
+    for link, expire_date in links:
+        response.append({"link": link, "expire_date": expire_date})
+    return response
 
-# @app.route('/<link>', methods=['GET'])
-# def link(link):
-#     redirect_link = log(link, flask.request)
-#     return flask.redirect(redirect_link)
 
+@app.post("/api/{link}/delete")
+async def link_delete(link: str, api_key: str = Security(check_api_key)):
+    """
+    Delete the specified link and all records associated with it
+    """
+    data = delete_link(link, api_key)
+    if data == 404:
+        raise HTTPException(
+            status_code=fastapi.status.HTTP_404_NOT_FOUND,
+            detail="Link does not exist"
+        )
+    if data == 401:
+        raise HTTPException(
+            status_code=fastapi.status.HTTP_401_UNAUTHORIZED,
+            detail="Link not associated with given API key"
+        )
+    else:
+        return {"link": f"The link {data} has been deleted"}
 
-# @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.post("/api/{link}/renew")
+async def link_renew(link: str, api_key: str = Security(check_api_key)):
+    """
+    Renew a specifiec link (adds 7 more days from the current date)
+    """
+    data = renew_link(link, api_key)
+    if data == 404:
+        raise HTTPException(
+            status_code=fastapi.status.HTTP_404_NOT_FOUND,
+            detail="Link does not exist"
+        )
+    if data == 401:
+        raise HTTPException(
+            status_code=fastapi.status.HTTP_401_UNAUTHORIZED,
+            detail="Link not associated with given API key"
+        )
+    else:
+        return {"link": f"The link {data[0]} has been renewed and will expire on {data[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.get("/api/{link}/records")
+async def link_records(link: str, api_key: str = Security(check_api_key)):
+    """
+    Retrieve all IP log records for the specified link
+    """
+    data = get_link_records(link, api_key)
+    if data == 404:
+        raise HTTPException(
+            status_code=fastapi.status.HTTP_404_NOT_FOUND,
+            detail="Link does not exist"
+        )
+    if data == 401:
+        raise HTTPException(
+            status_code=fastapi.status.HTTP_401_UNAUTHORIZED,
+            detail="Link not associated with given API key"
+        )
+    if data == 204:
+        raise HTTPException(
+            status_code=fastapi.status.HTTP_204_NO_CONTENT,
+            detail="No records found"
+        )
+    else:
+        response = []
+        for timestamp, ip, location, browser, os, user_agent, isp in data:
+            response.append({"timestamp": timestamp, "ip": ip, "location": location, "browser": browser, "os": os, "user_agent": user_agent, "isp": isp})
 
-# @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]
+        return response
 
 
-# @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.post("/api/{link}/delrecords")
+async def link_delrecords(link: str, api_key: str = Security(check_api_key)):
+    """
+    Delete all IP log records for the specified link
+    """
+    data = delete_link_records(link, api_key)
+    if data == 404:
+        raise HTTPException(
+            status_code=fastapi.status.HTTP_404_NOT_FOUND,
+            detail="Link does not exist"
+        )
+    if data == 401:
+        raise HTTPException(
+            status_code=fastapi.status.HTTP_401_UNAUTHORIZED,
+            detail="Link not associated with given API key"
+        )
+    else:
+        return {"link": f"The records for link {data} have been deleted"}
\ No newline at end of file