More secury (thanks �)

.github/workflows/tests.yml

@@ -28,5 +28,5 @@ jobs:
     
     - name: Start API server & run tests!
       run: |
+        python checks
         python run
-    

api/core.py

@@ -1,14 +1,25 @@
 """User management."""
 
 import os
-import json
+import sys
-import fastapi
 
-from db.users import UserManager
+project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+sys.path.append(project_root)
+
+# the code above is to allow importing from the root folder
+
+import os
+import json
+import hmac
+import fastapi
 
 from dhooks import Webhook, Embed
 from dotenv import load_dotenv
 
+import checks.client
+
+from db.users import UserManager
+
 load_dotenv()
 router = fastapi.APIRouter(tags=['core'])
 
@@ -21,11 +32,16 @@ async def check_core_auth(request):
     """
     received_auth = request.headers.get('Authorization')
 
-    if received_auth != os.getenv('CORE_API_KEY'):
+    correct_core_api = os.environ['CORE_API_KEY']
+
+    # use hmac.compare_digest to prevent timing attacks
+    if received_auth and hmac.compare_digest(received_auth, correct_core_api):
         return fastapi.Response(status_code=403, content='Invalid or no API key given.')
 
 @router.get('/users')
 async def get_users(discord_id: int, incoming_request: fastapi.Request):
+    """Returns a user by their discord ID. Requires a core API key."""
+
     auth = await check_core_auth(incoming_request)
     if auth:
         return auth
@@ -39,7 +55,9 @@ async def get_users(discord_id: int, incoming_request: fastapi.Request):
     return user
 
 async def new_user_webhook(user: dict) -> None:
-    dhook = Webhook(os.getenv('DISCORD_WEBHOOK__USER_CREATED'))
+    """Runs when a new user is created."""
+
+    dhook = Webhook(os.environ['DISCORD_WEBHOOK__USER_CREATED'])
 
     embed = Embed(
         description='New User',
@@ -54,6 +72,8 @@ async def new_user_webhook(user: dict) -> None:
 
 @router.post('/users')
 async def create_user(incoming_request: fastapi.Request):
+    """Creates a user. Requires a core API key."""
+
     auth_error = await check_core_auth(incoming_request)
 
     if auth_error:
@@ -74,6 +94,8 @@ async def create_user(incoming_request: fastapi.Request):
 
 @router.put('/users')
 async def update_user(incoming_request: fastapi.Request):
+    """Updates a user. Requires a core API key."""
+
     auth_error = await check_core_auth(incoming_request)
 
     if auth_error:
@@ -91,3 +113,19 @@ async def update_user(incoming_request: fastapi.Request):
     user = await manager.update_by_discord_id(discord_id, updates)
 
     return user
+
+@router.get('/checks')
+async def run_checks(incoming_request: fastapi.Request):
+    """Tests the API. Requires a core API key."""
+
+    auth_error = await check_core_auth(incoming_request)
+
+    if auth_error:
+        return auth_error
+
+    return {
+        'library': await checks.client.test_library(),
+        'library_moderation': await checks.client.test_library_moderation(),
+        'api_moderation': await checks.client.test_api_moderation(),
+        'models': await checks.client.test_models()
+    }

api/db/users.py

@@ -89,7 +89,7 @@ class UserManager:
         db = await self._get_collection('users')
         return await db.update_one({'_id': user_id}, update)
 
-    async def upate_by_discord_id(self, discord_id: str, update):
+    async def update_by_discord_id(self, discord_id: str, update):
         db = await self._get_collection('users')
         return await db.update_one({'auth.discord': str(int(discord_id))}, update)
 

api/main.py

@@ -24,21 +24,19 @@ app.include_router(core.router)
 
 @app.on_event('startup')
 async def startup_event():
-    # DATABASE FIX https://stackoverflow.com/questions/65970988/python-mongodb-motor-objectid-object-is-not-iterable-error-while-trying-to-f
+    """Runs when the API starts up."""
-    import pydantic, bson
-    # pydantic.json.ENCODERS_BY_TYPE[bson.objectid.ObjectId]=str
 
 @app.get('/')
 async def root():
     """
-    Returns the root endpoint.
+    Returns general information about the API.
     """
 
     return {
-        'status': 'ok',
+        'hi': 'Welcome to the Nova API!',
-        'usage_docs': 'https://nova-oss.com',
+        'learn_more_here': 'https://nova-oss.com',
-        'core_api_docs_for_developers': '/docs',
+        'github': 'https://github.com/novaoss/nova-api',
-        'github': 'https://github.com/novaoss/nova-api'
+        'core_api_docs_for_nova_developers': '/docs'
     }
 
 app.add_route('/{path:path}', transfer.handle, ['GET', 'POST', 'PUT', 'DELETE', 'PATCH'])

api/streaming.py

@@ -93,7 +93,7 @@ async def stream(
                     'cookies': incoming_request.cookies
                 })
         except ValueError as exc:
-            webhook = dhooks.Webhook(os.getenv('DISCORD_WEBHOOK__API_ISSUE'))
+            webhook = dhooks.Webhook(os.environ['DISCORD_WEBHOOK__API_ISSUE'])
             webhook.send(content=f'API Issue: **`{exc}`**\nhttps://i.imgflip.com/7uv122.jpg')
             yield await errors.yield_error(500, 'Sorry, the API has no working keys anymore.', 'The admins have been messaged automatically.')
             return

api/transfer.py

@@ -44,16 +44,16 @@ async def handle(incoming_request):
     received_key = incoming_request.headers.get('Authorization')
 
     if not received_key or not received_key.startswith('Bearer '):
-        return await errors.error(401, 'No NovaAI API key given!', 'Add "Authorization: Bearer nv-..." to your request headers.')
+        return await errors.error(403, 'No NovaAI API key given!', 'Add \'Authorization: Bearer nv-...\' to your request headers.')
 
     user = await users.user_by_api_key(received_key.split('Bearer ')[1].strip())
 
     if not user or not user['status']['active']:
-        return await errors.error(401, 'Invalid or inactive NovaAI API key!', 'Create a new NovaOSS API key or reactivate your account.')
+        return await errors.error(403, 'Invalid or inactive NovaAI API key!', 'Create a new NovaOSS API key or reactivate your account.')
 
     ban_reason = user['status']['ban_reason']
     if ban_reason:
-        return await errors.error(403, f'Your NovaAI account has been banned. Reason: "{ban_reason}".', 'Contact the staff for an appeal.')
+        return await errors.error(403, f'Your NovaAI account has been banned. Reason: \'{ban_reason}\'.', 'Contact the staff for an appeal.')
 
     costs = config['costs']
     cost = costs['other']

checks/__init__.py

@@ -0,0 +1 @@
+from . import client

checks/__main__.py

@@ -0,0 +1,2 @@
+import client
+client.demo()

checks/client.py

@@ -1,9 +1,11 @@
 """Tests the API."""
 
 import os
-import openai as closedai
-import httpx
 import time
+import httpx
+import openai
+import asyncio
+import traceback
 
 from rich import print
 from typing import List
@@ -22,7 +24,7 @@ MESSAGES = [
 
 api_endpoint = 'http://localhost:2332'
 
-def test_server():
+async def test_server():
     """Tests if the API server is running."""
 
     try:
@@ -30,7 +32,7 @@ def test_server():
     except httpx.ConnectError as exc:
         raise ConnectionError(f'API is not running on port {api_endpoint}.') from exc
 
-def test_api(model: str=MODEL, messages: List[dict]=None) -> dict:
+async def test_api(model: str=MODEL, messages: List[dict]=None) -> dict:
     """Tests an API api_endpoint."""
 
     json_data = {
@@ -49,10 +51,10 @@ def test_api(model: str=MODEL, messages: List[dict]=None) -> dict:
 
     return response.text
 
-def test_library():
+async def test_library():
-    """Tests if the api_endpoint is working with the Python library."""
+    """Tests if the api_endpoint is working with the OpenAI Python library."""
 
-    completion = closedai.ChatCompletion.create(
+    completion = openai.ChatCompletion.create(
         model=MODEL,
         messages=MESSAGES
     )
@@ -61,13 +63,13 @@ def test_library():
 
     return completion['choices'][0]['message']['content']
 
-def test_library_moderation():
+async def test_library_moderation():
     try:
-        return closedai.Moderation.create('I wanna kill myself, I wanna kill myself; It\'s all I hear right now, it\'s all I hear right now')
+        return openai.Moderation.create('I wanna kill myself, I wanna kill myself; It\'s all I hear right now, it\'s all I hear right now')
-    except closedai.error.InvalidRequestError:
+    except openai.error.InvalidRequestError:
         return True
 
-def test_models():
+async def test_models():
     response = httpx.get(
         url=f'{api_endpoint}/models',
         headers=HEADERS,
@@ -76,7 +78,7 @@ def test_models():
     response.raise_for_status()
     return response.json()
 
-def test_api_moderation() -> dict:
+async def test_api_moderation() -> dict:
     """Tests an API api_endpoint."""
 
     response = httpx.get(
@@ -90,38 +92,43 @@ def test_api_moderation() -> dict:
 
 # ==========================================================================================
 
-def test_all():
+def demo():
     """Runs all tests."""
-    try:
-        print("Waiting until API Server is started up...")
-        time.sleep(6)
 
-        print('[lightblue]Running test on API server to check if its running...')
+    try:
-        print(test_server())
+        for _ in range(30):
+            if test_server():
+                break
+
+            print('Waiting until API Server is started up...')
+            time.sleep(1)
+        else:
+            raise ConnectionError('API Server is not running.')
 
         print('[lightblue]Running a api endpoint to see if requests can go through...')
-        print(test_api('gpt-3.5-trubo'))
+        print(asyncio.run(test_api('gpt-3.5-turbo')))
 
         print('[lightblue]Checking if the API works with the python library...')
-        print(test_library())
+        print(asyncio.run(test_library()))
 
         print('[lightblue]Checking if the moderation endpoint works...')
-        print(test_library_moderation())
+        print(asyncio.run(test_library_moderation()))
 
         print('[lightblue]Checking the /v1/models endpoint...')
-        print(test_models())
+        print(asyncio.run(test_models()))
-    except Exception as e:
+
-        print('[red]Error: ')
+    except Exception as exc:
-        print(e)
+        print('[red]Error: ' + str(exc))
+        traceback.print_exc()
         exit(500)
 
-if __name__ == '__main__':
+openai.api_base = api_endpoint
-    closedai.api_base = api_endpoint
+openai.api_key = os.environ['NOVA_KEY']
-    closedai.api_key = os.environ['NOVA_KEY']
 
 HEADERS = {
     'Content-Type': 'application/json',
-        'Authorization': 'Bearer ' + closedai.api_key
+    'Authorization': 'Bearer ' + openai.api_key
 }
 
-    test_all()
+if __name__ == '__main__':
+    demo()

screen.sh

@@ -12,5 +12,8 @@ cp env/.prod.env /home/nova-prod/.env
 # Change directory
 cd /home/nova-prod
 
+# Kill the production server
+fuser -k 2333/tcp
+
 # Start screen
 screen -S nova-api python run prod

setup.md

@@ -100,7 +100,7 @@ python run 1337
 ```
 
 ## Test if it works
-`python tests`
+`python checks`
 
 ## Ports
 ```yml