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)