Functional simple user management

README.md

@@ -1,6 +1,39 @@
-# ☄️ Nova API Server
+# ☄️ NovaOSS API Server
 Reverse proxy server for "Closed"AI's API.
 
+## NovaOSS APIs
+Our infrastructure might seem a bit confusing, but it's actually quite simple. Just the first one really matters for you, if you want to access our AI API. The other ones are just for the team.
+
+### AI API
+**Public** (everyone can use it with a valid API key)
+
+Official endpoints: `https://api.nova-oss.com/v1/...`
+Documentation & info: [nova-oss.com](https://nova-oss.com)
+
+- Access to AI models
+
+***
+
+### User/Account management API
+**Private** (NovaOSS operators only!)
+
+Official endpoints: `https://api.nova-oss.com/...`
+Documentation: [api.nova-oss.com/docs](https://api.nova-oss.com/docs)
+
+- Access to user accounts
+- Implemented in [NovaCord](https://nova-oss.com/novacord)
+
+### Website API
+**Private** (NovaOSS operators only!)
+
+Official endpoints: `https://nova-oss.com/api/...`
+
+This one's code can be found in the following repository: [github.com/novaoss/nova-web](https://github.com/novaoss/nova-web)
+
+- Used for the Terms of Service (ToS) verification for the Discord bot.
+- In a different repository and with a different domain because it needs to display codes on the website.
+- Implemented in [NovaCord](https://nova-oss.com/novacord)
+
 ## Install
 Assuming you have a new version of Python 3 and pip installed:
 ```py
@@ -26,6 +59,10 @@ pip install .
 ```
 
 ## `.env` configuration
+Create a `.env` file, make sure not to reveal it to anyone, and fill in the required values in the format `KEY=VALUE`. Otherwise, the code won't run.
+
+### Database
+- `API_DB_PATH` the path to the databases, e.g. `/etc/nova/db/.` (this way, the database `users` would be saved in `/etc/nova/db/.users.db`.)
 
 ### Proxy
 - `PROXY_TYPE` (optional, defaults to `socks.PROXY_TYPE_HTTP`): the type of proxy - can be `http`, `https`, `socks4`, `socks5`, `4` or `5`, etc... 
@@ -41,14 +78,53 @@ pip install .
 ### `ACTUAL_IPS` (optional)
 This is a security measure to make sure a proxy, VPN, Tor or any other IP hiding service is used by the host when accessing "Closed"AI's API.
 It is a space separated list of IP addresses that are allowed to access the API.
-You can also just add the *beginning* of an API address, like `12.123.` to allow all IPs starting with `12.123.`.
+You can also just add the *beginning* of an API address, like `12.123.` (without an asterisk!) to allow all IPs starting with `12.123.`.
+> To disable the warning if you don't have this feature enabled, set `ACTUAL_IPS` to `None`.
 
 ### `DEMO_AUTH`
 API key for demo purposes. You can give this to trusted team members. Never use it in production.
 
-> To disable the warning if you don't have this feature enabled, set `ACTUAL_IPS` to any value.
+### `CORE_API_KEY`
+This will 
+
 ## Run
-`python run`
+> **Warning:** read the according section for production usage!
+
+For developement:
+
+```bash
+python run
+```
+
+This will run the development server on port `2332`.
+
+You can also specify a port, e.g.:
+
+```bash
+python run 1337
+```
 
 ## Test if it works
 `python tests`
+
+## Ports
+```yml
+2332: Developement (default)
+2333: Production
+```
+
+## Production
+Make sure your server is secure and up to date.
+Check everything.
+
+The following command will run the API  __without__ a reloader!
+
+```bash
+python run prod
+```
+
+or 
+
+```bash
+./screen.sh
+```

api/apihandler.py

@@ -1,13 +0,0 @@
-from typing import Union, Optional
-
-class Request:
-    def __init__(self,
-        method: str,
-        url: str,
-        json_payload: Optional[Union[dict, list]]=None,
-        headers: dict=None
-    ):
-        self.method = method
-        self.url = url
-        self.json = json_payload
-        self.headers = headers or {}

api/closed_keys.py

@@ -1,33 +1,17 @@
+"""ClosedAI key manager."""
+
 import os
-import sys
 import uuid
 import time
 import asyncio
-import aiosqlite
 
-from dotenv import load_dotenv
+from rich import print
-
+from helpers import databases
-load_dotenv()
-
-db_uri = os.getenv('DB_PATH')
-
-try:
-    os.remove(db_uri)
-except FileNotFoundError:
-    pass
-
-async def to_dict(cursor, row):
-    return {col[0]: row[idx] for idx, col in enumerate(cursor.description)}
-
-async def connect_db():
-    """Creates a connection to the database"""
-
-    return await aiosqlite.connect(db_uri)
 
 async def prepare() -> None:
     """Creates the database tables"""
 
-    keys_db = await connect_db()
+    keys_db = await databases.connect('closed_keys')
     await keys_db.execute(
         """
         CREATE TABLE IF NOT EXISTS closed_keys (
@@ -55,38 +39,29 @@ async def add_key(
 
     tags = tags or []
 
-    keys_db = await connect_db()
+    new_key = {
-
+        'id': str(uuid.uuid4()),
-    parameters = {
+        'key': key,
-        "id": str(uuid.uuid4()),
+        'source': source,
-        "key": key,
+        'created_at': int(time.time()),
-        "source": source,
+        'last_used': -1,
-        "created_at": int(time.time()),
+        'uses_count': 0,
-        "last_used": -1,
+        'tokens_generated': 0,
-        "uses_count": 0,
+        'active': True,
-        "tokens_generated": 0,
+        'working': True,
-        "active": True,
+        'tags': '/'.join(tags),
-        "working": True,
-        "tags": '/'.join(tags),
     }
 
-    sep = ', '
+    await databases.insert_dict(new_key, 'closed_keys', 'closed_keys')
-    query = f"""
-        INSERT INTO closed_keys ({sep.join(parameters.keys())})
-        VALUES ({sep.join([f':{key}' for key in parameters.keys()])})
-    """
-
-    await keys_db.execute(query, parameters)
-    await keys_db.commit()
 
 async def get_working_key() -> dict:
     """Returns a working key"""
 
-    keys_db = await connect_db()
+    keys_db = await databases.connect('closed_keys')
 
     async with keys_db.execute('SELECT * FROM closed_keys WHERE working = 1') as cursor:
         async for row in cursor:
-            return await to_dict(cursor, row)
+            return await databases.row_to_dict(row, cursor)
 
     return None
 
@@ -95,7 +70,7 @@ asyncio.run(prepare())
 async def key_stopped_working(key: str) -> None:
     """Marks a key as stopped working"""
 
-    keys_db = await connect_db()
+    keys_db = await databases.connect('closed_keys')
 
     await keys_db.execute(
         """
@@ -110,7 +85,7 @@ async def key_stopped_working(key: str) -> None:
 async def key_was_used(key: str, num_tokens: int) -> None:
     """Updates the stats of a key"""
 
-    keys_db = await connect_db()
+    keys_db = await databases.connect('closed_keys')
 
     # set last_used to int of time.time(), adds one to uses_count and adds num_tokens to tokens_generated
     await keys_db.execute(

api/core.py

@@ -0,0 +1,48 @@
+"""User management."""
+
+import os
+import json
+import fastapi
+
+import users
+
+from dotenv import load_dotenv
+
+load_dotenv()
+router = fastapi.APIRouter(tags=['core'])
+
+async def check_core_auth(request):
+    received_auth = request.headers.get('Authorization')
+
+    if received_auth != os.getenv('CORE_API_KEY'):
+        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):
+    auth_error = await check_core_auth(incoming_request)
+
+    if auth_error:
+        return auth_error
+
+    user = await users.get_user(by_discord_id=discord_id)
+
+    if not user:
+        return fastapi.Response(status_code=404, content='User not found.')
+    return user
+
+@router.post('/users')
+async def create_user(incoming_request: fastapi.Request):
+    auth_error = await check_core_auth(incoming_request)
+
+    if auth_error:
+        return auth_error
+
+    try:
+        payload = await incoming_request.json()
+        discord_id = payload.get('discord_id')
+    except (json.decoder.JSONDecodeError, AttributeError):
+        return fastapi.Response(status_code=400, content='Invalid or no payload received.')
+
+    user = await users.add_user(discord_id=discord_id)
+
+    return user

api/helpers/databases.py

@@ -0,0 +1,36 @@
+"""Database helper."""
+
+import os
+import aiosqlite
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+async def row_to_dict(row, cursor):
+    """Converts a database row to into a <dict>."""    
+
+    return {col[0]: row[idx] for idx, col in enumerate(cursor.description)}
+
+async def connect(name: str):
+    """Creates a connection to the database"""
+
+    return await aiosqlite.connect(f'{os.getenv("API_DB_PATH")}{name}.db')
+
+async def insert_dict(dict_: dict, table: str, db) -> dict:
+    """Adds a dictionary to a table, safely."""
+
+    if isinstance(db, str):
+        db = await connect(db)
+
+    sep = ', '
+    query = f"""
+        INSERT INTO {table} ({sep.join(dict_.keys())})
+        VALUES ({sep.join([f':{key}' for key in dict_.keys()])})
+    """
+
+    await db.execute(query, dict_)
+    await db.commit()
+    
+    return dict_
+

api/helpers/requesting.py

@@ -1,3 +1,5 @@
+"""Manages web requests."""
+
 import os
 
 from dotenv import load_dotenv
@@ -26,3 +28,8 @@ class Request:
         self.payload = payload
         self.headers = headers
         self.timeout = int(os.getenv('TRANSFER_TIMEOUT', '120'))
+
+class HTTPXRequest(Request):
+    def __init__(self, url: str, *args, **kwargs):
+        super().__init__(url, *args, **kwargs)
+        self.url += '?httpx=1'

api/main.py

@@ -1,12 +1,13 @@
 """FastAPI setup."""
 
 import fastapi
-import asyncio
 
 from fastapi.middleware.cors import CORSMiddleware
 
 from dotenv import load_dotenv
 
+import core
+import users
 import transfer
 
 load_dotenv()
@@ -21,9 +22,11 @@ app.add_middleware(
     allow_headers=['*']
 )
 
+app.include_router(core.router)
+
 @app.on_event('startup')
 async def startup_event():
-    """Read up the API server."""
+    await users.prepare()
 
 @app.get('/')
 async def root():
@@ -31,15 +34,9 @@ async def root():
 
     return {
         'status': 'ok',
-        'readme': 'https://nova-oss.com'
+        'usage_docs': 'https://nova-oss.com',
+        'core_api_docs_for_developers': '/docs',
+        'github': 'https://github.com/novaoss/nova-api'
     }
 
-@app.route('/v1')
+app.add_route('/v1/{path:path}', transfer.handle_api_request, ['GET', 'POST', 'PUT', 'DELETE', 'PATCH'])
-async def api_root():
-    """Returns the API root endpoint."""
-
-    return {
-        'status': 'ok',
-    }
-
-app.add_route('/{path:path}', transfer.handle_api_request, ['GET', 'POST', 'PUT', 'DELETE', 'PATCH'])

api/netclient.py

@@ -4,7 +4,7 @@ import httpx
 import proxies
 
 from dotenv import load_dotenv
-from request_manager import Request
+from helpers.requesting import Request
 
 load_dotenv()
 

api/proxies.py

@@ -81,16 +81,6 @@ default_proxy = Proxy(
     password=os.getenv('PROXY_PASS')
 )
 
-def test_requests():
-    import requests
-
-    # return requests.get(
-    #     'https://checkip.amazonaws.com',
-    #     timeout=5,
-    #     proxies=default_proxy.urls
-    # ).text.strip()
-
-
 def test_httpx():
     import httpx
 
@@ -102,5 +92,4 @@ def test_httpx():
     ).text.strip()
 
 if __name__ == '__main__':
-    print(test_requests())
     print(test_httpx())

api/transfer.py

@@ -6,9 +6,9 @@ import logging
 import starlette
 
 import netclient
-import request_manager
 
 from dotenv import load_dotenv
+from helpers import requesting
 
 load_dotenv()
 
@@ -40,12 +40,10 @@ async def handle_api_request(incoming_request, target_endpoint: str=''):
     except json.decoder.JSONDecodeError:
         payload = {}
 
-    target_provider = 'moe'
-
     if 'temperature' in payload or 'functions' in payload:
         target_provider = 'closed'
 
-    request = request_manager.Request(
+    request = requesting.Request(
         url=target_url,
         payload=payload,
         method=incoming_request.method,

api/users.py

@@ -0,0 +1,115 @@
+"""User system."""
+
+import os
+import uuid
+import time
+import string
+import random
+import asyncio
+
+from helpers import databases
+from dotenv import load_dotenv
+
+load_dotenv()
+
+async def prepare() -> None:
+    """Creates the database tables"""
+
+    users_db = await databases.connect('users')
+    await users_db.execute(
+        """
+        CREATE TABLE IF NOT EXISTS users (
+            id TEXT PRIMARY KEY,
+            api_key TEXT,
+            active BOOLEAN,
+            created_at INTEGER,
+            last_used INTEGER DEFAULT 0,
+            uses_count INTEGER DEFAULT 0,
+            tokens_generated INTEGER DEFAULT 0,
+            discord_id INTEGER DEFAULT 0,
+            credit INTEGER DEFAULT 0,
+            tags TEXT DEFAULT ''
+        )
+        """
+    )
+    await users_db.commit()
+
+async def add_user(
+    discord_id: int=0,
+    tags: list=None,
+) -> dict:
+    """Adds a new key to the database"""
+
+    chars = string.ascii_letters + string.digits
+
+    infix = os.getenv('KEYGEN_INFIX')
+    suffix = ''.join(random.choices(chars, k=20))
+    prefix = ''.join(random.choices(chars, k=20))
+
+    key = f'nv-{prefix}{infix}{suffix}'
+
+    tags = tags or []    
+    new_user = {
+        'id': str(uuid.uuid4()),
+        'api_key': key,
+        'created_at': int(time.time()),
+        'last_used': 0,
+        'uses_count': 0,
+        'tokens_generated': 0,
+        'active': True,
+        'discord_id': discord_id,
+        'credit': 0,
+        'tags': '/'.join(tags)
+    }
+
+    await databases.insert_dict(new_user, 'users', 'users')
+    return new_user
+
+async def get_user(
+    by_id: str='',
+    by_discord_id: int=0,
+):
+    users_db = await databases.connect('users')
+
+    async with users_db.execute(
+        'SELECT * FROM users WHERE id = :id OR discord_id = :discord_id',
+        {'id': by_id, 'discord_id': by_discord_id}
+    ) as cursor:
+        async for row in cursor:
+            result = await databases.row_to_dict(row, cursor)
+            return result
+
+    return None
+
+async def get_all_users():
+    users_db = await databases.connect('users')
+    results = []
+
+    async with users_db.execute(
+        'SELECT * FROM users'
+    ) as cursor:
+        async for row in cursor:
+            result = await databases.row_to_dict(row, cursor)
+            results.append(result)
+
+    return results
+
+async def demo():
+    await prepare()
+
+    users = await get_all_users()
+    print(users)
+
+    example_id = 133769420
+    user = await add_user(discord_id=example_id)
+    print(user)
+
+    del user
+    print('Fetching user...')
+
+    user = await get_user(by_discord_id=example_id) 
+    print(user['api_key'])
+
+if __name__ == '__main__':
+    asyncio.run(demo())
+    os.system(f'pkill -f {os.path.basename(__file__)}')

run/__main__.py

@@ -1,5 +1,25 @@
-import sys
+"""Starts the API.
-import os
 
-port = sys.argv[1] if len(sys.argv) > 1 else 2333
+Usage:
-os.system(f'cd api && uvicorn main:app --reload --host 0.0.0.0 --port {port}')
+$ python run 1234
+Runs on port 1234.
+
+$ python run prod
+Runs for production.
+
+$ python run 1234 prod
+Runs for production on the speicified port.
+
+"""
+
+import os
+import sys
+
+port = sys.argv[1] if len(sys.argv) > 1 else 2332
+dev = True
+
+if 'prod' in sys.argv:
+    port = 2333
+    dev = False
+
+os.system(f'cd api && uvicorn main:app{" --reload" if dev else ""} --host 0.0.0.0 --port {port}')

screen.sh

@@ -0,0 +1 @@
+screen -s nova-api python run prod

tests/__main__.py

@@ -21,7 +21,7 @@ MESSAGES = [
     },
 ]
 
-api_endpoint = 'http://localhost:2333'
+api_endpoint = 'http://localhost:2332'
 
 def test_server():
     """Tests if the API server is running."""