From cdcb1eedd1a5cd6b2071ebea776e52ab7100fc4f Mon Sep 17 00:00:00 2001 From: nsde Date: Tue, 1 Aug 2023 02:38:55 +0200 Subject: [PATCH] Functional simple user management --- README.md | 84 ++++++++++++- api/apihandler.py | 13 -- api/{closedkeys.py => closed_keys.py} | 67 ++++------ api/core.py | 48 ++++++++ api/helpers/databases.py | 36 ++++++ .../requesting.py} | 7 ++ api/main.py | 21 ++-- api/netclient.py | 2 +- api/proxies.py | 11 -- api/transfer.py | 6 +- api/users.py | 115 ++++++++++++++++++ run/__main__.py | 28 ++++- screen.sh | 1 + tests/__main__.py | 2 +- 14 files changed, 345 insertions(+), 96 deletions(-) delete mode 100644 api/apihandler.py rename api/{closedkeys.py => closed_keys.py} (65%) create mode 100644 api/core.py create mode 100644 api/helpers/databases.py rename api/{request_manager.py => helpers/requesting.py} (77%) create mode 100644 api/users.py create mode 100755 screen.sh diff --git a/README.md b/README.md index 01b8dfd..b2b774b 100644 --- a/README.md +++ b/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 +``` \ No newline at end of file diff --git a/api/apihandler.py b/api/apihandler.py deleted file mode 100644 index 0f099ae..0000000 --- a/api/apihandler.py +++ /dev/null @@ -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 {} diff --git a/api/closedkeys.py b/api/closed_keys.py similarity index 65% rename from api/closedkeys.py rename to api/closed_keys.py index 1604283..b492227 100644 --- a/api/closedkeys.py +++ b/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 - -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) +from rich import print +from helpers import databases 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() - - parameters = { - "id": str(uuid.uuid4()), - "key": key, - "source": source, - "created_at": int(time.time()), - "last_used": -1, - "uses_count": 0, - "tokens_generated": 0, - "active": True, - "working": True, - "tags": '/'.join(tags), + new_key = { + 'id': str(uuid.uuid4()), + 'key': key, + 'source': source, + 'created_at': int(time.time()), + 'last_used': -1, + 'uses_count': 0, + 'tokens_generated': 0, + 'active': True, + 'working': True, + 'tags': '/'.join(tags), } - sep = ', ' - 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() + await databases.insert_dict(new_key, 'closed_keys', 'closed_keys') 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( diff --git a/api/core.py b/api/core.py new file mode 100644 index 0000000..74f291f --- /dev/null +++ b/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 diff --git a/api/helpers/databases.py b/api/helpers/databases.py new file mode 100644 index 0000000..b8c2a56 --- /dev/null +++ b/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 .""" + + 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_ + diff --git a/api/request_manager.py b/api/helpers/requesting.py similarity index 77% rename from api/request_manager.py rename to api/helpers/requesting.py index b503bfb..e4cc49c 100644 --- a/api/request_manager.py +++ b/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' diff --git a/api/main.py b/api/main.py index 5397b6b..23b808e 100644 --- a/api/main.py +++ b/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') -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']) +app.add_route('/v1/{path:path}', transfer.handle_api_request, ['GET', 'POST', 'PUT', 'DELETE', 'PATCH']) diff --git a/api/netclient.py b/api/netclient.py index 00ffebe..3de4653 100644 --- a/api/netclient.py +++ b/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() diff --git a/api/proxies.py b/api/proxies.py index 9c2cacf..ff022d3 100644 --- a/api/proxies.py +++ b/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()) diff --git a/api/transfer.py b/api/transfer.py index 5f617d9..cb813f7 100644 --- a/api/transfer.py +++ b/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, diff --git a/api/users.py b/api/users.py new file mode 100644 index 0000000..5235be7 --- /dev/null +++ b/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__)}') diff --git a/run/__main__.py b/run/__main__.py index c1dff99..ced6c7d 100644 --- a/run/__main__.py +++ b/run/__main__.py @@ -1,5 +1,25 @@ -import sys -import os +"""Starts the API. -port = sys.argv[1] if len(sys.argv) > 1 else 2333 -os.system(f'cd api && uvicorn main:app --reload --host 0.0.0.0 --port {port}') +Usage: +$ 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}') diff --git a/screen.sh b/screen.sh new file mode 100755 index 0000000..18923f7 --- /dev/null +++ b/screen.sh @@ -0,0 +1 @@ +screen -s nova-api python run prod diff --git a/tests/__main__.py b/tests/__main__.py index 4c1e9de..230bf66 100644 --- a/tests/__main__.py +++ b/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."""