2023-08-04 03:30:56 +02:00
|
|
|
import os
|
|
|
|
import pytz
|
|
|
|
import asyncio
|
|
|
|
import datetime
|
|
|
|
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
from motor.motor_asyncio import AsyncIOMotorClient
|
|
|
|
|
|
|
|
load_dotenv()
|
|
|
|
|
2023-08-12 18:20:18 +02:00
|
|
|
## Statistics
|
2023-08-04 03:30:56 +02:00
|
|
|
|
2023-08-14 05:11:15 +02:00
|
|
|
class StatsManager:
|
2023-08-13 18:42:38 +02:00
|
|
|
"""
|
|
|
|
### The manager for all statistics tracking
|
|
|
|
Stats tracked:
|
|
|
|
- Dates
|
|
|
|
- IPs
|
|
|
|
- Target URLs
|
|
|
|
- Tokens
|
|
|
|
- Models
|
|
|
|
- URL Paths
|
|
|
|
"""
|
2023-08-14 05:11:15 +02:00
|
|
|
|
|
|
|
def __init__(self):
|
2023-08-19 16:26:14 +02:00
|
|
|
self.conn = AsyncIOMotorClient(os.environ['MONGO_URI'])
|
2023-08-14 05:11:15 +02:00
|
|
|
|
|
|
|
async def _get_collection(self, collection_name: str):
|
2023-08-17 13:13:20 +02:00
|
|
|
return self.conn[os.getenv('MONGO_NAME', 'nova-test')][collection_name]
|
2023-08-14 05:11:15 +02:00
|
|
|
|
|
|
|
async def add_date(self):
|
2023-08-13 18:42:38 +02:00
|
|
|
date = datetime.datetime.now(pytz.timezone('GMT')).strftime('%Y.%m.%d')
|
|
|
|
year, month, day = date.split('.')
|
2023-08-04 03:30:56 +02:00
|
|
|
|
2023-08-14 05:11:15 +02:00
|
|
|
db = await self._get_collection('stats')
|
2023-08-13 18:42:38 +02:00
|
|
|
await db.update_one({}, {'$inc': {f'dates.{year}.{month}.{day}': 1}}, upsert=True)
|
2023-08-04 03:30:56 +02:00
|
|
|
|
2023-08-14 05:11:15 +02:00
|
|
|
async def add_ip_address(self, ip_address: str):
|
2023-08-13 18:42:38 +02:00
|
|
|
ip_address = ip_address.replace('.', '_')
|
2023-08-14 05:11:15 +02:00
|
|
|
db = await self._get_collection('stats')
|
2023-08-13 18:42:38 +02:00
|
|
|
await db.update_one({}, {'$inc': {f'ips.{ip_address}': 1}}, upsert=True)
|
2023-08-04 03:30:56 +02:00
|
|
|
|
2023-08-14 05:11:15 +02:00
|
|
|
async def add_target(self, url: str):
|
|
|
|
db = await self._get_collection('stats')
|
2023-08-13 18:42:38 +02:00
|
|
|
await db.update_one({}, {'$inc': {f'targets.{url}': 1}}, upsert=True)
|
2023-08-04 03:30:56 +02:00
|
|
|
|
2023-08-14 05:11:15 +02:00
|
|
|
async def add_tokens(self, tokens: int, model: str):
|
|
|
|
db = await self._get_collection('stats')
|
2023-08-13 18:42:38 +02:00
|
|
|
await db.update_one({}, {'$inc': {f'tokens.{model}': tokens}}, upsert=True)
|
2023-08-04 03:30:56 +02:00
|
|
|
|
2023-08-14 05:11:15 +02:00
|
|
|
async def add_model(self, model: str):
|
|
|
|
db = await self._get_collection('stats')
|
2023-08-13 18:42:38 +02:00
|
|
|
await db.update_one({}, {'$inc': {f'models.{model}': 1}}, upsert=True)
|
2023-08-04 03:30:56 +02:00
|
|
|
|
2023-08-14 05:11:15 +02:00
|
|
|
async def add_path(self, path: str):
|
2023-08-13 18:42:38 +02:00
|
|
|
path = path.replace('/', '_')
|
2023-08-14 05:11:15 +02:00
|
|
|
db = await self._get_collection('stats')
|
2023-08-13 18:42:38 +02:00
|
|
|
await db.update_one({}, {'$inc': {f'paths.{path}': 1}}, upsert=True)
|
2023-08-04 03:30:56 +02:00
|
|
|
|
2023-08-14 05:11:15 +02:00
|
|
|
async def get_value(self, obj_filter):
|
|
|
|
db = await self._get_collection('stats')
|
2023-08-13 18:42:38 +02:00
|
|
|
return await db.find_one({obj_filter})
|
2023-08-04 03:30:56 +02:00
|
|
|
|
2023-08-28 00:58:32 +02:00
|
|
|
manager = StatsManager()
|
|
|
|
|
2023-08-04 03:30:56 +02:00
|
|
|
if __name__ == '__main__':
|
2023-08-17 13:11:35 +02:00
|
|
|
stats = StatsManager()
|
|
|
|
asyncio.run(stats.add_date())
|
|
|
|
asyncio.run(stats.add_path('/__demo/test'))
|