feat: add support for optional cookie in API requests and enhance command context handling

Author: eWloYW8

CSOJ-cli.py

@@ -6,7 +6,7 @@
 import sys
 from pathlib import Path
 import base64
-import fnmatch  # 导入用于文件名模式匹配的库
+import fnmatch
 
 import click
 import requests
@@ -79,10 +79,13 @@ def load_config():
         sys.exit(1)
 
 
-def get_api_headers():
+def get_api_headers(cookie=None):
     """Gets the necessary headers for authenticated API calls."""
     config = load_config()
-    return {"Authorization": f"Bearer {config['jwt']}"}
+    headers = {"Authorization": f"Bearer {config['jwt']}"}
+    if cookie:
+        headers['Cookie'] = cookie
+    return headers
 
 
 def handle_api_error(response):
@@ -98,13 +101,17 @@ def handle_api_error(response):
 
 # --- CLI Command Group ---
 @click.group(cls=AliasedGroup)
-def cli():
+@click.option('--cookie', help='Optional cookie string to include in all requests.')
+@click.pass_context
+def cli(ctx, cookie):
     """
     CSOJ-cli: A command-line tool for the CSOJ Online Judge platform.
     
     Supports both full command names and short aliases (e.g., 'leaderboard' or 'lb').
+    The --cookie option can be used with any command.
     """
-    pass
+    ctx.ensure_object(dict)
+    ctx.obj['cookie'] = cookie
 
 
 # --- CLI Commands ---
@@ -124,14 +131,17 @@ def login(domain, jwt):
 
 
 @cli.command("ls-contests", help="List all available contests. (alias: lsc)")
-def list_contests():
+@click.pass_context
+def list_contests(ctx):
     config = load_config()
     domain = config["domain"]
+    cookie = ctx.obj.get('cookie')
+    headers = get_api_headers(cookie)
     url = f"{domain}/api/v1/contests"
 
     try:
         with console.status("[bold green]Fetching contests...[/]"):
-            response = requests.get(url, timeout=10)
+            response = requests.get(url, headers=headers, timeout=10)
 
         if response.status_code != 200:
             handle_api_error(response)
@@ -165,14 +175,17 @@ def list_contests():
 
 @cli.command("ls-problems", help="List all problems for a given contest. (alias: lsp)")
 @click.argument("contest_id")
-def list_problems(contest_id):
+@click.pass_context
+def list_problems(ctx, contest_id):
     config = load_config()
     domain = config["domain"]
+    cookie = ctx.obj.get('cookie')
+    headers = get_api_headers(cookie)
     contest_url = f"{domain}/api/v1/contests/{contest_id}"
 
     try:
         with console.status("[bold green]Fetching contest details...[/]"):
-            contest_response = requests.get(contest_url, timeout=10)
+            contest_response = requests.get(contest_url, headers=headers, timeout=10)
 
         if contest_response.status_code != 200:
             handle_api_error(contest_response)
@@ -200,7 +213,7 @@ def list_problems(contest_id):
             task = progress.add_task("[green]Fetching problems...", total=len(problem_ids))
             for prob_id in problem_ids:
                 problem_url = f"{domain}/api/v1/problems/{prob_id}"
-                problem_response = requests.get(problem_url, timeout=10)
+                problem_response = requests.get(problem_url, headers=headers, timeout=10)
                 if problem_response.status_code == 200:
                     problem_data = problem_response.json().get("data", {})
                     table.add_row(problem_data.get("id", "N/A"), problem_data.get("name", "N/A"))
@@ -216,10 +229,12 @@ def list_problems(contest_id):
 
 @cli.command(help="Register for a specific contest. (alias: reg)")
 @click.argument("contest_id")
-def register(contest_id):
+@click.pass_context
+def register(ctx, contest_id):
     config = load_config()
     domain = config["domain"]
-    headers = get_api_headers()
+    cookie = ctx.obj.get('cookie')
+    headers = get_api_headers(cookie)
     url = f"{domain}/api/v1/contests/{contest_id}/register"
 
     try:
@@ -241,16 +256,18 @@ def register(contest_id):
 @cli.command(help="Submit a file or folder to a problem. (alias: sub)")
 @click.argument("problem_id")
 @click.argument("path", type=click.Path(exists=True, resolve_path=True, file_okay=True, dir_okay=True))
-def submit(problem_id, path):
+@click.pass_context
+def submit(ctx, problem_id, path):
     config = load_config()
     domain = config["domain"]
-    headers = get_api_headers()
+    cookie = ctx.obj.get('cookie')
+    headers = get_api_headers(cookie)
     submit_url = f"{domain}/api/v1/problems/{problem_id}/submit"
     problem_url = f"{domain}/api/v1/problems/{problem_id}"
     submit_path = Path(path)
-    files_to_upload = []
     skipped_files = []
-    prepared_files = []
+    files_to_upload_info = []
+    prepared_files_for_request = []
 
     try:
         # Step 1: Fetch problem details to get upload rules
@@ -275,7 +292,7 @@ def submit(problem_id, path):
 
         for file_path in all_files_to_check:
             if submit_path.is_dir():
-                relative_path = file_path.relative_to(submit_path)
+                relative_path = file_path.relative_to(submit_path.parent)
             else:
                 relative_path = Path(file_path.name)
 
@@ -286,7 +303,7 @@ def submit(problem_id, path):
                 is_allowed = any(fnmatch.fnmatch(str(relative_path), pattern) for pattern in upload_rules)
 
             if is_allowed:
-                files_to_upload.append({'absolute_path': file_path, 'relative_path': relative_path})
+                files_to_upload_info.append({'absolute_path': file_path, 'relative_path': relative_path})
             else:
                 skipped_files.append(str(relative_path))
 
@@ -296,27 +313,26 @@ def submit(problem_id, path):
             for skipped in skipped_files:
                 console.print(f"  - [dim]{skipped}[/dim]")
 
-        if not files_to_upload:
+        if not files_to_upload_info:
             console.print("[bold red]Error:[/bold red] No valid files found for submission after filtering. Aborting.")
             sys.exit(1)
 
         # Step 4: Prepare files for upload by opening them
-        for file_info in files_to_upload:
+        for file_info in files_to_upload_info:
             file_obj = open(file_info['absolute_path'], 'rb')
-            b64_name = base64.b64encode(str(file_info['relative_path']).encode()).decode()
-            prepared_files.append(('files', (b64_name, file_obj)))
+            b64_name = base64.b64encode(file_info['relative_path'].as_posix().encode()).decode()
+            prepared_files_for_request.append(('files', (b64_name, file_obj)))
 
-        console.print(f"Uploading {len(prepared_files)} file(s) to problem [cyan]{problem_id}[/cyan]...")
-        response = requests.post(submit_url, headers=headers, files=prepared_files, timeout=60)
+        console.print(f"Uploading {len(prepared_files_for_request)} file(s) to problem [cyan]{problem_id}[/cyan]...")
+        response = requests.post(submit_url, headers=headers, files=prepared_files_for_request, timeout=60)
 
         if response.status_code == 403:
-            # Check for ban message first
             try:
                 error_data = response.json()
                 if "banned" in error_data.get("message", "").lower():
-                    handle_api_error(response) # Let the handler print the specific ban message
+                    handle_api_error(response)
             except json.JSONDecodeError:
-                pass # Fallback to generic registration message
+                pass
 
             console.print("[bold red]Submission Forbidden (HTTP 403).[/bold red]")
             console.print("This may be because you have not registered for the contest.")
@@ -340,12 +356,12 @@ def submit(problem_id, path):
         sys.exit(1)
     finally:
         # Step 5: Ensure all opened file objects are closed
-        for _, (_, file_obj) in prepared_files:
+        for _, (_, file_obj) in prepared_files_for_request:
             if file_obj:
                 file_obj.close()
 
 
-async def stream_logs(domain, jwt, submission_id, container_id, step_name):
+async def stream_logs(domain, jwt, submission_id, container_id, step_name, cookie=None):
     """Coroutine to connect to WebSocket and stream logs."""
     ws_protocol = "wss" if domain.startswith("https") else "ws"
     http_protocol = "https" if ws_protocol == "wss" else "http"
@@ -358,8 +374,12 @@ async def stream_logs(domain, jwt, submission_id, container_id, step_name):
     connecting_text.overflow = "fold"
     console.print(Panel(connecting_text, title=panel_title, border_style="green", title_align="left"))
 
+    ws_headers = {}
+    if cookie:
+        ws_headers['Cookie'] = cookie
+
     try:
-        async with websockets.connect(uri) as websocket:
+        async with websockets.connect(uri, extra_headers=ws_headers) as websocket:
             async for message in websocket:
                 try:
                     log_data = json.loads(message)
@@ -388,10 +408,12 @@ async def stream_logs(domain, jwt, submission_id, container_id, step_name):
 @cli.command(help="Get the status and optionally logs of a submission. (alias: st)")
 @click.argument("submission_id")
 @click.option("--logs", "-l", is_flag=True, help="Stream logs for all workflow steps.")
-def status(submission_id, logs):
+@click.pass_context
+def status(ctx, submission_id, logs):
     config = load_config()
     domain = config["domain"]
-    headers = get_api_headers()
+    cookie = ctx.obj.get('cookie')
+    headers = get_api_headers(cookie)
     sub_url = f"{domain}/api/v1/submissions/{submission_id}"
 
     try:
@@ -444,7 +466,7 @@ def status(submission_id, logs):
                     console.print(Panel(f"[dim]Logs for step '[blue]{step_name}[/blue]' are hidden by the problem setter.[/dim]", border_style="dim"))
                     continue
 
-                asyncio.run(stream_logs(domain, config['jwt'], submission_id, container.get("id"), step_name))
+                asyncio.run(stream_logs(domain, config['jwt'], submission_id, container.get("id"), step_name, cookie))
 
     except requests.exceptions.RequestException as e:
         console.print(f"[bold red]Network Error:[/bold red] {e}")
@@ -453,14 +475,17 @@ def status(submission_id, logs):
 
 @cli.command(help="Display the leaderboard for a contest. (alias: lb)")
 @click.argument("contest_id")
-def leaderboard(contest_id):
+@click.pass_context
+def leaderboard(ctx, contest_id):
     config = load_config()
     domain = config["domain"]
+    cookie = ctx.obj.get('cookie')
+    headers = get_api_headers(cookie)
 
     try:
         with console.status("[bold green]Fetching contest details...[/]"):
             contest_url = f"{domain}/api/v1/contests/{contest_id}"
-            contest_res = requests.get(contest_url, timeout=10)
+            contest_res = requests.get(contest_url, headers=headers, timeout=10)
             if contest_res.status_code != 200:
                 handle_api_error(contest_res)
             contest_data = contest_res.json().get("data", {})
@@ -469,7 +494,7 @@ def leaderboard(contest_id):
 
         with console.status(f"[bold green]Fetching leaderboard for {contest_name}...[/]"):
             leaderboard_url = f"{domain}/api/v1/contests/{contest_id}/leaderboard"
-            leaderboard_res = requests.get(leaderboard_url, timeout=15)
+            leaderboard_res = requests.get(leaderboard_url, headers=headers, timeout=15)
             if leaderboard_res.status_code != 200:
                 handle_api_error(leaderboard_res)
             leaderboard_data = leaderboard_res.json().get("data", [])
@@ -507,14 +532,17 @@ def leaderboard(contest_id):
 
 @cli.command("show-problem", help="Show the description and details of a problem. (alias: show)")
 @click.argument("problem_id")
-def show_problem(problem_id):
+@click.pass_context
+def show_problem(ctx, problem_id):
     config = load_config()
     domain = config["domain"]
+    cookie = ctx.obj.get('cookie')
+    headers = get_api_headers(cookie)
     url = f"{domain}/api/v1/problems/{problem_id}"
 
     try:
         with console.status(f"[bold green]Fetching problem {problem_id}...[/]"):
-            response = requests.get(url, timeout=10)
+            response = requests.get(url, headers=headers, timeout=10)
 
         if response.status_code != 200:
             handle_api_error(response)
@@ -538,10 +566,12 @@ def show_problem(problem_id):
 
 
 @cli.command("ls-submissions", help="List your recent submissions. (alias: lss)")
-def list_submissions():
+@click.pass_context
+def list_submissions(ctx):
     config = load_config()
     domain = config["domain"]
-    headers = get_api_headers()
+    cookie = ctx.obj.get('cookie')
+    headers = get_api_headers(cookie)
     url = f"{domain}/api/v1/submissions"
 
     try: