diff --git a/dumpyarabot/handlers.py b/dumpyarabot/handlers.py index 4262b81..deb204d 100644 --- a/dumpyarabot/handlers.py +++ b/dumpyarabot/handlers.py @@ -20,10 +20,14 @@ async def dump( message: Optional[Message] = update.effective_message if not chat or not message: + console.print("[red]Chat or message object is None[/red]") return # Ensure it can only be used in the correct group if chat.id not in settings.ALLOWED_CHATS: + console.print( + f"[yellow]Unauthorized chat attempt from chat_id: {chat.id}[/yellow]" + ) await context.bot.send_message( chat_id=chat.id, reply_to_message_id=message.message_id, @@ -33,6 +37,7 @@ async def dump( # Ensure that we had some arguments passed if not context.args: + console.print("[yellow]No arguments provided for dump command[/yellow]") usage = "Usage: `/dump [URL] [a|f|b|p]`\nURL: required, a: alt dumper, f: force, b: blacklist, p: use privdump" await context.bot.send_message( chat_id=chat.id, @@ -48,9 +53,27 @@ async def dump( add_blacklist = "b" in context.args[1:] if len(context.args) > 1 else False use_privdump = "p" in context.args[1:] if len(context.args) > 1 else False + console.print(f"[green]Dump request:[/green]") + console.print(f" URL: {url}") + console.print(f" Alt dumper: {use_alt_dumper}") + console.print(f" Force: {force}") + console.print(f" Blacklist: {add_blacklist}") + console.print(f" Privdump: {use_privdump}") + # Delete the user's message immediately if privdump is used if use_privdump: - await context.bot.delete_message(chat_id=chat.id, message_id=message.message_id) + console.print( + f"[blue]Privdump requested - deleting message {message.message_id}[/blue]" + ) + try: + await context.bot.delete_message( + chat_id=chat.id, message_id=message.message_id + ) + console.print( + "[green]Successfully deleted original message for privdump[/green]" + ) + except Exception as e: + console.print(f"[red]Failed to delete message for privdump: {e}[/red]") # Try to check for existing build and call jenkins if necessary try: @@ -62,6 +85,7 @@ async def dump( ) if not force: + console.print("[blue]Checking for existing builds...[/blue]") initial_message = await context.bot.send_message( chat_id=chat.id, reply_to_message_id=None if use_privdump else message.message_id, @@ -70,6 +94,9 @@ async def dump( exists, status_message = await utils.check_existing_build(dump_args) if exists: + console.print( + f"[yellow]Found existing build: {status_message}[/yellow]" + ) await context.bot.edit_message_text( chat_id=chat.id, message_id=initial_message.message_id, @@ -82,12 +109,17 @@ async def dump( message_id=initial_message.message_id, ) + console.print("[blue]Calling Jenkins to start build...[/blue]") response_text = await utils.call_jenkins(dump_args) + console.print(f"[green]Jenkins response: {response_text}[/green]") + except ValidationError: + console.print(f"[red]Invalid URL provided: {url}[/red]") response_text = "Invalid URL" except Exception: + console.print("[red]Unexpected error occurred:[/red]") + console.print_exception() response_text = "An error occurred" - console.print_exception(show_locals=True) # Reply to the user with whatever the status is await context.bot.send_message( @@ -104,10 +136,14 @@ async def cancel_dump(update: Update, context: ContextTypes.DEFAULT_TYPE) -> Non user = update.effective_user if not chat or not message or not user: + console.print("[red]Chat, message or user object is None[/red]") return # Ensure it can only be used in the correct group if chat.id not in settings.ALLOWED_CHATS: + console.print( + f"[yellow]Unauthorized chat attempt for cancel from chat_id: {chat.id}[/yellow]" + ) await context.bot.send_message( chat_id=chat.id, reply_to_message_id=message.message_id, @@ -118,6 +154,9 @@ async def cancel_dump(update: Update, context: ContextTypes.DEFAULT_TYPE) -> Non # Check if the user is an admin admins = await chat.get_administrators() if user not in [admin.user for admin in admins]: + console.print( + f"[yellow]Non-admin user {user.id} tried to use cancel command[/yellow]" + ) await context.bot.send_message( chat_id=chat.id, reply_to_message_id=message.message_id, @@ -127,6 +166,7 @@ async def cancel_dump(update: Update, context: ContextTypes.DEFAULT_TYPE) -> Non # Ensure that we had some arguments passed if not context.args: + console.print("[yellow]No job_id provided for cancel command[/yellow]") usage = ( "Usage: `/cancel [job_id] [p]`\njob_id: required, p: cancel privdump job" ) @@ -141,7 +181,21 @@ async def cancel_dump(update: Update, context: ContextTypes.DEFAULT_TYPE) -> Non job_id = context.args[0] use_privdump = "p" in context.args[1:] if len(context.args) > 1 else False - response_message = await utils.cancel_jenkins_job(job_id, use_privdump) + console.print(f"[blue]Cancel request:[/blue]") + console.print(f" Job ID: {job_id}") + console.print(f" Privdump: {use_privdump}") + console.print(f" Requested by: {user.username} (ID: {user.id})") + + try: + response_message = await utils.cancel_jenkins_job(job_id, use_privdump) + console.print( + f"[green]Successfully processed cancel request: {response_message}[/green]" + ) + except Exception as e: + console.print("[red]Error processing cancel request:[/red]") + console.print_exception() + response_message = f"Error cancelling job: {str(e)}" + await context.bot.send_message( chat_id=chat.id, reply_to_message_id=message.message_id, diff --git a/dumpyarabot/utils.py b/dumpyarabot/utils.py index df0258a..cfcc1b6 100644 --- a/dumpyarabot/utils.py +++ b/dumpyarabot/utils.py @@ -1,107 +1,163 @@ from typing import List, Tuple import httpx +from rich.console import Console from dumpyarabot import schemas from dumpyarabot.config import settings +console = Console() + async def get_jenkins_builds(job_name: str) -> List[schemas.JenkinsBuild]: """Fetch all builds from Jenkins for a specific job.""" + console.print(f"[blue]Fetching builds for job: {job_name}[/blue]") async with httpx.AsyncClient() as client: - response = await client.get( - f"{settings.JENKINS_URL}/job/{job_name}/api/json", - params={"tree": "allBuilds[number,result,actions[parameters[name,value]]]"}, - auth=(settings.JENKINS_USER_NAME, settings.JENKINS_USER_TOKEN), - ) - response.raise_for_status() - return [schemas.JenkinsBuild(**build) for build in response.json()["allBuilds"]] - - -async def check_existing_build(args: schemas.DumpArguments) -> Tuple[bool, str]: - """Check if a build with the given parameters already exists.""" - job_name = "privdump" if args.use_privdump else "dumpyara" - builds = await get_jenkins_builds(job_name) - - for build in builds: - if _is_matching_build(build, args): - return _get_build_status(build) - - return False, f"No matching build found. A new {job_name} build will be started." + try: + response = await client.get( + f"{settings.JENKINS_URL}/job/{job_name}/api/json", + params={ + "tree": "allBuilds[number,result,actions[parameters[name,value]]]" + }, + auth=(settings.JENKINS_USER_NAME, settings.JENKINS_USER_TOKEN), + ) + response.raise_for_status() + builds = [ + schemas.JenkinsBuild(**build) for build in response.json()["allBuilds"] + ] + console.print( + f"[green]Successfully fetched {len(builds)} builds for {job_name}[/green]" + ) + return builds + except Exception as e: + console.print(f"[red]Failed to fetch builds for {job_name}: {e}[/red]") + raise def _is_matching_build( build: schemas.JenkinsBuild, args: schemas.DumpArguments ) -> bool: """Check if a build matches the given arguments.""" + console.print("[blue]Checking build parameters match...[/blue]") for action in build.actions: if "parameters" in action: params = {param["name"]: param["value"] for param in action["parameters"]} - return ( + matches = ( params.get("URL") == args.url.unicode_string() and params.get("USE_ALT_DUMPER") == args.use_alt_dumper and params.get("ADD_BLACKLIST") == args.add_blacklist ) + if matches: + console.print("[green]Found matching build parameters[/green]") + return matches return False def _get_build_status(build: schemas.JenkinsBuild) -> Tuple[bool, str]: """Get the status of a build.""" + console.print(f"[blue]Checking build status: #{build.number}[/blue]") if build.result is None: + console.print("[yellow]Build is currently in progress[/yellow]") return ( True, f"Build #{build.number} is currently in progress for this URL and settings.", ) elif build.result == "SUCCESS": + console.print("[green]Build completed successfully[/green]") return ( True, f"Build #{build.number} has already successfully completed for this URL and settings.", ) else: + console.print( + f"[yellow]Build result was {build.result}, will start new build[/yellow]" + ) return ( False, f"Build #{build.number} exists for this URL and settings, but result was {build.result}. A new build will be started.", ) +async def check_existing_build(args: schemas.DumpArguments) -> Tuple[bool, str]: + """Check if a build with the given parameters already exists.""" + job_name = "privdump" if args.use_privdump else "dumpyara" + console.print(f"[blue]Checking existing builds for {job_name}[/blue]") + console.print("Build parameters:", args) + + builds = await get_jenkins_builds(job_name) + + for build in builds: + if _is_matching_build(build, args): + status = _get_build_status(build) + console.print(f"[yellow]Found matching build - Status: {status}[/yellow]") + return status + + console.print(f"[green]No matching build found for {job_name}[/green]") + return False, f"No matching build found. A new {job_name} build will be started." + + async def call_jenkins(args: schemas.DumpArguments) -> str: """Call Jenkins to start a new build.""" job_name = "privdump" if args.use_privdump else "dumpyara" + console.print(f"[blue]Starting new {job_name} build[/blue]") + console.print("Build parameters:", args) + async with httpx.AsyncClient() as client: - response = await client.post( - f"{settings.JENKINS_URL}/job/{job_name}/buildWithParameters", - params={ - "URL": args.url.unicode_string(), - "USE_ALT_DUMPER": args.use_alt_dumper, - "ADD_BLACKLIST": args.add_blacklist, - }, - auth=(settings.JENKINS_USER_NAME, settings.JENKINS_USER_TOKEN), - ) - response.raise_for_status() - return f"{job_name.capitalize()} job started" + try: + response = await client.post( + f"{settings.JENKINS_URL}/job/{job_name}/buildWithParameters", + params={ + "URL": args.url.unicode_string(), + "USE_ALT_DUMPER": args.use_alt_dumper, + "ADD_BLACKLIST": args.add_blacklist, + }, + auth=(settings.JENKINS_USER_NAME, settings.JENKINS_USER_TOKEN), + ) + response.raise_for_status() + console.print(f"[green]Successfully started {job_name} build[/green]") + return f"{job_name.capitalize()} job started" + except Exception as e: + console.print(f"[red]Failed to start {job_name} build: {e}[/red]") + raise async def cancel_jenkins_job(job_id: str, use_privdump: bool = False) -> str: """Cancel a Jenkins job.""" job_name = "privdump" if use_privdump else "dumpyara" + console.print(f"[blue]Attempting to cancel {job_name} job {job_id}[/blue]") + async with httpx.AsyncClient() as client: - response = await client.post( - f"{settings.JENKINS_URL}/job/{job_name}/{job_id}/stop", - auth=(settings.JENKINS_USER_NAME, settings.JENKINS_USER_TOKEN), - follow_redirects=True, - ) - if response.status_code == 200: - return f"Job with ID {job_id} has been cancelled in {job_name}." - elif response.status_code == 404: + try: response = await client.post( - f"{settings.JENKINS_URL}/queue/cancelItem", - params={"id": job_id}, + f"{settings.JENKINS_URL}/job/{job_name}/{job_id}/stop", auth=(settings.JENKINS_USER_NAME, settings.JENKINS_USER_TOKEN), follow_redirects=True, ) - if response.status_code == 204: - return ( - f"Job with ID {job_id} has been removed from the {job_name} queue." + if response.status_code == 200: + console.print( + f"[green]Successfully cancelled {job_name} job {job_id}[/green]" ) - - return f"Failed to cancel job with ID {job_id} in {job_name}. Job not found or already completed." + return f"Job with ID {job_id} has been cancelled in {job_name}." + elif response.status_code == 404: + console.print( + f"[yellow]Job {job_id} not found, checking queue[/yellow]" + ) + response = await client.post( + f"{settings.JENKINS_URL}/queue/cancelItem", + params={"id": job_id}, + auth=(settings.JENKINS_USER_NAME, settings.JENKINS_USER_TOKEN), + follow_redirects=True, + ) + if response.status_code == 204: + console.print( + f"[green]Successfully removed {job_name} job {job_id} from queue[/green]" + ) + return f"Job with ID {job_id} has been removed from the {job_name} queue." + + console.print(f"[yellow]Failed to cancel {job_name} job {job_id}[/yellow]") + return f"Failed to cancel job with ID {job_id} in {job_name}. Job not found or already completed." + except Exception as e: + console.print( + f"[red]Error while cancelling {job_name} job {job_id}: {e}[/red]" + ) + raise