246 lines
8.6 KiB
Python
246 lines
8.6 KiB
Python
import argparse
|
|
import asyncio
|
|
import json
|
|
from aiohttp import web
|
|
from base64 import b64decode
|
|
from typing import Dict, Any
|
|
import uuid
|
|
|
|
import dotenv
|
|
import os
|
|
|
|
import re
|
|
import logging
|
|
from logging.handlers import RotatingFileHandler
|
|
from pathlib import Path
|
|
|
|
print("Creating log directory if it doesn't exist...")
|
|
os.makedirs("./logs/", exist_ok=True)
|
|
|
|
import banner
|
|
import get_sq_env
|
|
from time import time
|
|
|
|
time = int(time())
|
|
import asquare as sq
|
|
|
|
# from asquareclass import AsyncSquareClient # type: ignore
|
|
from abricksetapi import AsyncBricksetAPI
|
|
|
|
print("Loading environment variables...")
|
|
dotenv.load_dotenv()
|
|
SQUARE_LEGO_CAT_ID = os.environ.get("SQUARE_LEGO_CAT_ID")
|
|
|
|
from asyncio_throttle.throttler import Throttler
|
|
|
|
# set up logger
|
|
print("Setting up logger...")
|
|
logger = logging.getLogger("studspider")
|
|
logger.setLevel(logging.DEBUG)
|
|
handler = RotatingFileHandler("./logs/studspider.log", maxBytes=125000, backupCount=5)
|
|
logger.addHandler(handler)
|
|
|
|
lego = AsyncBricksetAPI(debug=True, logger=logger)
|
|
# sq = AsyncSquareClient(logger=logger)
|
|
# logging.basicConfig(
|
|
# filename="./logs/studspider.log", encoding="utf-8", level=logging.DEBUG
|
|
# )
|
|
print("Setting up rate limiter...")
|
|
print("Doing RE things...")
|
|
UA = re.compile(r"^kopling\/?", re.IGNORECASE)
|
|
|
|
|
|
async def run_update():
|
|
set_list = await sq.acat2id(SQUARE_LEGO_CAT_ID) # type: ignore
|
|
# print(f"[@DEBUG]\t {set_list}") # type: ignore
|
|
for l_set in set_list:
|
|
# print(f"[@DEBUG]{l_set}")
|
|
try:
|
|
l_setinfo = await lego.throttled_get_set(l_set["number"])
|
|
print(f"[@DEBUG] Found {l_setinfo["pw_name"]}")
|
|
except Exception as e:
|
|
print(f"Error fetching set\t: {l_set}: {str(e)}")
|
|
continue
|
|
|
|
# print(f"[@DEBUG]{l_setinfo}")
|
|
# await sq.update_itm(l_set, l_setinfo)
|
|
"""for l_set in set_list:
|
|
print(f"[@DEBUG]{l_set}")
|
|
try:
|
|
l_setinfo = await lego.throttled_get_set(l_set["number"])
|
|
except Exception as e:
|
|
print(f"Error fetching set\t: {l_set}: {str(e)}")
|
|
continue
|
|
|
|
print(f"[@DEBUG]{l_setinfo}")"""
|
|
# await sq.update_itm(l_set, l_setinfo)
|
|
|
|
|
|
def valid_uuid(val):
|
|
try:
|
|
uuid.UUID(str(val))
|
|
return True
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
async def hook_handler(request: web.Request) -> web.Response:
|
|
"""Handle incoming JSON webhook POSTs on /hook endpoint"""
|
|
try:
|
|
# Verify request method
|
|
if request.method != "POST":
|
|
logger.error(
|
|
f"{time}\tInvalid request method to /hook from %s", request.remote
|
|
)
|
|
return web.Response(status=405, text="Method not allowed. Use POST.")
|
|
|
|
# Get JSON data from request
|
|
try:
|
|
data = await request.json()
|
|
except json.JSONDecodeError:
|
|
logger.error("Invalid JSON to /hook from %s", request.remote)
|
|
return web.Response(status=400, text="Invalid JSON in request body")
|
|
|
|
# Log the received webhook
|
|
print(f"[@DEBUG] Webhook received at {time} from {request.remote}")
|
|
logger.info(
|
|
f"{time}\tWebhook received from {request.remote}: {json.dumps(data)}"
|
|
)
|
|
|
|
# Process the webhook data here
|
|
# TODO: Add your webhook processing logic
|
|
|
|
response_data = {
|
|
"status": "success",
|
|
"message": "Webhook received",
|
|
"timestamp": time,
|
|
"data_received": data,
|
|
}
|
|
|
|
return web.json_response(response_data)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Webhook handler error: {str(e)}")
|
|
print(f"Webhook handler error: {str(e)}")
|
|
return web.Response(status=500, text="Internal server error")
|
|
|
|
|
|
async def request_handler(request: web.Request) -> web.Response:
|
|
"""Handle incoming POST requests with command processing"""
|
|
try:
|
|
# Verify request method
|
|
if request.method != "POST":
|
|
logger.error(f"{time}\tInvalid request method from %s at ", request.remote)
|
|
return web.Response(status=405, text="Method not allowed. Use POST.")
|
|
|
|
# Get JSON data from request
|
|
try:
|
|
data = await request.json()
|
|
except json.JSONDecodeError:
|
|
logger.error("Invalid JSON from %s", request.remote)
|
|
return web.Response(
|
|
status=400, text="{time()}\tInvalid JSON in request body"
|
|
)
|
|
# Validate Idempotency-Key header
|
|
if valid_uuid(request.headers.get("Idempotency-Key")) is False:
|
|
logger.error("{time()}\tInvalid Idempotency-Key from %s", request.remote)
|
|
return web.Response(status=412, text="412 Headers rejected!")
|
|
# Validate User-Agent header
|
|
if not re.match(UA, str(request.headers.get("User-Agent"))):
|
|
logger.error("{time()}\tInvalid User-Agent from %s", request.remote)
|
|
return web.Response(status=412, text="412 Headers rejected!")
|
|
# Validate command exists
|
|
if "CMD" not in data:
|
|
logger.error("{time()}\tMissing CMD field from %s", request.remote)
|
|
return web.Response(status=400, text="Missing CMD field in request")
|
|
|
|
# Decode and process command
|
|
try:
|
|
decoded_cmd = b64decode(data["CMD"]).decode("utf-8")
|
|
print(
|
|
f"""At: {time}\nReceived command: {decoded_cmd}\nUser-Agent: {request.headers.get('User-Agent')}\nIdempotency-Key: {request.headers.get('Idempotency-Key')}"""
|
|
)
|
|
logging.info(
|
|
f"{time}\tCommand: {decoded_cmd}\tFrom: {request.remote}\tUA: {request.headers.get('User-Agent')}\tIdempotency-Key: {request.headers.get('Idempotency-Key')}"
|
|
)
|
|
|
|
# Process the command
|
|
if decoded_cmd == "RUNUPDATE":
|
|
response_data = {"status": "success", "message": "Update started"}
|
|
await run_update()
|
|
logging.info(f"{time()}\tUpdate command processed successfully")
|
|
return web.json_response(response_data)
|
|
else:
|
|
return web.Response(status=400, text=f"Unknown command: {decoded_cmd}")
|
|
|
|
except Exception as e:
|
|
print(f"Error processing command: {str(e)}")
|
|
return web.Response(status=500, text=f"Server error: {str(e)}")
|
|
|
|
except Exception as e:
|
|
print(f"Request handler error: {str(e)}")
|
|
return web.Response(status=500, text="Internal server error")
|
|
|
|
|
|
print("Starting web server...")
|
|
os.system("cls" if os.name == "nt" else "clear")
|
|
app = web.Application()
|
|
app.router.add_post("/api", request_handler)
|
|
app.router.add_post("/hook", hook_handler)
|
|
|
|
parser = argparse.ArgumentParser(description="studspider")
|
|
parser.add_argument(
|
|
"--host",
|
|
type=str,
|
|
default="localhost",
|
|
help="IP address to listen on. default: localhost",
|
|
)
|
|
parser.add_argument(
|
|
"--port", type=int, default=7432, help="Port to listen on. default: 7432"
|
|
)
|
|
parser.add_argument(
|
|
"--hook-port",
|
|
type=int,
|
|
default=8080,
|
|
help="Optional separate port for webhook endpoint. default: 8080",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if __name__ == "__main__":
|
|
print(banner.generate("studspider v0.1"))
|
|
print("Built for PLAY-WELL TEKNOLOGIES by contact@traps.moe\n\n")
|
|
|
|
# If hook-port is specified, run a separate server for webhooks
|
|
if args.hook_port:
|
|
print(f"Starting main server on {args.host}:{args.port}")
|
|
print(f"Starting webhook server on {args.host}:{args.hook_port}")
|
|
|
|
# Create separate app for webhook
|
|
hook_app = web.Application()
|
|
hook_app.router.add_post("/hook", hook_handler)
|
|
|
|
# Run both servers concurrently
|
|
async def run_servers():
|
|
runner1 = web.AppRunner(app)
|
|
await runner1.setup()
|
|
site1 = web.TCPSite(runner1, args.host, args.port)
|
|
|
|
runner2 = web.AppRunner(hook_app)
|
|
await runner2.setup()
|
|
site2 = web.TCPSite(runner2, args.host, args.hook_port)
|
|
|
|
await site1.start()
|
|
await site2.start()
|
|
|
|
print(f"Main server running on http://{args.host}:{args.port}")
|
|
print(f"Webhook server running on http://{args.host}:{args.hook_port}/hook")
|
|
|
|
# Keep running
|
|
await asyncio.Event().wait()
|
|
|
|
asyncio.run(run_servers())
|
|
else:
|
|
print(f"Server running on {args.host}:{args.port}")
|
|
print(f"Webhook endpoint: http://{args.host}:{args.port}/hook")
|
|
web.run_app(app, host=args.host, port=args.port)
|