feat: enhance login process with CSRF token

Author: yuvrajsinh5252

src/commands/login.py

@@ -5,22 +5,27 @@
 
 def login():
     """Login to LeetCode"""
-    # First try to use saved session
     if auth_manager.is_authenticated:
         saved_session = auth_manager.session_manager.load_session()
         if saved_session:
             typer.echo(typer.style(f"Already logged in as {saved_session['user_name']}!", fg=typer.colors.GREEN))
             return
 
-    typer.echo(typer.style("To login, you'll need your LEETCODE_SESSION token:", fg=typer.colors.YELLOW))
+    typer.echo(typer.style("To login, you'll need both CSRF and LEETCODE_SESSION tokens:", fg=typer.colors.YELLOW))
     typer.echo("\n1. Open LeetCode in your browser")
     typer.echo("2. Press F12 to open Developer Tools")
     typer.echo("3. Go to Application tab > Cookies > leetcode.com")
-    typer.echo("4. Find and copy the 'LEETCODE_SESSION' value\n")
+    typer.echo("4. Find and copy both 'csrftoken' and 'LEETCODE_SESSION' values\n")
 
-    leetcode_session = typer.prompt("Please enter your LEETCODE_SESSION token")
+    csrf_token = typer.prompt("Please enter your CSRF token")
+    csrf_result = auth_manager.verify_csrf_token(csrf_token)
+
+    if not csrf_result["success"]:
+        typer.echo(typer.style(f"\n✗ CSRF verification failed: {csrf_result['message']}", fg=typer.colors.RED))
+        return
 
-    result = auth_manager.login_with_session(leetcode_session)
+    leetcode_session = typer.prompt("Please enter your LEETCODE_SESSION token")
+    result = auth_manager.login_with_session(csrf_token, leetcode_session)
 
     if result["success"]:
         typer.echo(typer.style(f"\n✓ Successfully logged in as {result['user_name']}!", fg=typer.colors.GREEN))

src/server/auth.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, Any
 import requests
 from .session_manager import SessionManager
 
@@ -14,21 +14,56 @@ def _load_saved_session(self):
         """Try to load and validate saved session"""
         saved_session = self.session_manager.load_session()
         if saved_session:
-            result = self.login_with_session(saved_session['session_token'])
+            result = self.login_with_session(saved_session["csrftoken"], saved_session['session_token'])
             return result["success"]
         return False
 
-    def login_with_session(self, leetcode_session: str) -> Dict[str, any]: # type: ignore
-        """
-        Login to LeetCode using LEETCODE_SESSION token.
-        Returns a dictionary with success status and message.
-        """
+    def verify_csrf_token(self, csrf_token: str) -> Dict[str, Any]:
+        """Verify CSRF token by making a request to LeetCode's GraphQL endpoint"""
         try:
+            if not csrf_token:
+                return {"success": False, "message": "CSRF token is required"}
+
+            self.session.cookies.set('csrftoken', csrf_token, domain='leetcode.com')
+
+            response = self.session.post(
+                f"{self.BASE_URL}/graphql",
+                json={
+                    "query": "query { userStatus { isSignedIn username } }"
+                },
+                headers={
+                    'X-CSRFToken': csrf_token,
+                    'Accept': 'application/json',
+                    'Content-Type': 'application/json',
+                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
+                }
+            )
+
+            self.session.cookies.clear(domain='leetcode.com')
+
+            if response.status_code == 200:
+                data = response.json()
+                if 'errors' not in data:
+                    return {"success": True, "message": "CSRF token verified"}
+
+            return {"success": False, "message": "Invalid CSRF token"}
+
+        except Exception as e:
+            return {"success": False, "message": f"CSRF verification error: {str(e)}"}
+
+    def login_with_session(self, csrf_token: str, leetcode_session: str) -> Dict[str, Any]:
+        """Login using verified CSRF token and session token"""
+        try:
+
+            if not csrf_token or not leetcode_session:
+                return {"success": False, "message": "Both CSRF and LEETCODE_SESSION tokens are required"}
+
             self.session.cookies.set('LEETCODE_SESSION', leetcode_session, domain='leetcode.com')
 
             response = self.session.get(
                 f"{self.BASE_URL}/api/problems/all/",
                 headers={
+                    'X-CSRFToken': csrf_token,
                     'Accept': 'application/json',
                     'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
                 }
@@ -38,25 +73,15 @@ def login_with_session(self, leetcode_session: str) -> Dict[str, any]: # type: i
                 user_data = response.json()
                 if user_data.get('user_name'):
                     self.is_authenticated = True
-                    # Save the valid session
-                    self.session_manager.save_session(leetcode_session, user_data['user_name'])
+                    self.session_manager.save_session(csrf_token, leetcode_session, user_data['user_name'])
                     return {
                         "success": True,
                         "message": "Successfully logged in",
                         "user_name": user_data['user_name']
                     }
-                else:
-                    self.session_manager.clear_session()
-                    return {
-                        "success": False,
-                        "message": "Invalid session token"
-                    }
-            else:
-                self.session_manager.clear_session()
-                return {
-                    "success": False,
-                    "message": f"Login failed with status code: {response.status_code}"
-                }
+
+            self.session_manager.clear_session()
+            return {"success": False, "message": "Invalid session credentials"}
 
         except Exception as e:
             self.session_manager.clear_session()

src/server/session_manager.py

@@ -14,9 +14,10 @@ def _ensure_config_dir(self):
         """Ensure the config directory exists"""
         self.config_dir.mkdir(parents=True, exist_ok=True)
 
-    def save_session(self, session_token: str, user_name: str):
+    def save_session(self, csrftoken: str, session_token: str, user_name: str):
         """Save the session token and username to file"""
         data = {
+            'csrftoken': csrftoken,
             'session_token': session_token,
             'user_name': user_name
         }

src/server/solution_manager.py

@@ -15,10 +15,7 @@ def _clean_session_cookies(self):
         if not hasattr(self.session, 'cookies'):
             return
 
-        # Get all cookies
         all_cookies = list(self.session.cookies)
-
-        # Clear all cookies
         self.session.cookies.clear()
 
         # Add back only the most recent cookie for each name
@@ -32,7 +29,7 @@ def _get_csrf_token(self):
         for cookie in self.session.cookies:
             if cookie.name == 'csrftoken':
                 return cookie.value
-        return "iIf1V3zVGiU5F0neqw63pFm0YlFtk9i531xUBoQe0hZ06pmDGPZ0uJW6vyhr8GEH"
+        return ''
 
     def get_question_data(self, question_identifier: str) -> Dict[str, Any]:
         """Get question details using GraphQL
@@ -83,28 +80,58 @@ def get_question_data(self, question_identifier: str) -> Dict[str, Any]:
 
     def submit_solution(self, title_slug: str, code: str, lang: str = "python3") -> Dict[str, Any]:
         """Submit a solution to LeetCode"""
+        try:
+            self._clean_session_cookies()
 
-        # First get the question ID
-        question_data = self.get_question_data(title_slug)
-        question_id = question_data['data']['question']['questionId']
+            # Get question ID
+            question_data = self.get_question_data(title_slug)
+            question_data = question_data.get('data', {}).get('question')
+            if not question_data:
+                return {"success": False, "error": "Question data not found"}
 
-        submit_url = f"{self.BASE_URL}/problems/{title_slug}/submit/"
+            question_id = question_data['questionId']
+            submit_url = f"{self.BASE_URL}/problems/{title_slug}/submit/"
 
-        data = {
-            "lang": lang,
-            "question_id": question_id,
-            "typed_code": code
-        }
+            csrf_token = self._get_csrf_token()
+            typer.echo(f"Using CSRF token: {csrf_token}")
 
-        response = self.session.post(submit_url, json=data)
-        if response.status_code != 200:
-            return {"success": False, "error": f"Submission failed: {response.status_code}"}
+            headers = {
+                'referer': f"{self.BASE_URL}/problems/{title_slug}/",
+                'content-type': 'application/json',
+                'x-csrftoken': csrf_token,
+                'x-requested-with': 'XMLHttpRequest',
+                'origin': self.BASE_URL
+            }
 
-        submission_id = response.json().get('submission_id')
-        if not submission_id:
-            return {"success": False, "error": "No submission ID received"}
+            data = {
+                "lang": lang,
+                "question_id": question_id,
+                "typed_code": code
+            }
+
+            typer.echo(f"Sending request to {submit_url}")
+            response = self.session.post(submit_url, json=data, headers=headers)
+
+            if response.status_code != 200:
+                typer.echo(f"Error response: {response.text}", err=True)
+                return {"success": False, "error": f"Submission failed with status {response.status_code}"}
+
+            try:
+                result_data = response.json()
+                submission_id = result_data.get('submission_id')
+                if submission_id:
+                    typer.echo(f"Got submission ID: {submission_id}")
+                    return self.get_submission_result(submission_id)
+                else:
+                    typer.echo("No submission ID in response", err=True)
+                    return {"success": False, "error": "No submission ID received"}
+            except ValueError as e:
+                typer.echo(f"Failed to parse response: {response.text}", err=True)
+                return {"success": False, "error": f"Failed to parse response: {str(e)}"}
 
-        return self.get_submission_result(submission_id)
+        except Exception as e:
+            typer.echo(f"Submission error: {str(e)}", err=True)
+            return {"success": False, "error": f"Submission error: {str(e)}"}
 
     def test_solution(self, title_slug: str, code: str, lang: str = "python3", full: bool = False) -> Dict[str, Any]:
         """Test a solution with LeetCode test cases"""
@@ -157,7 +184,7 @@ def test_solution(self, title_slug: str, code: str, lang: str = "python3", full:
                 submission_id = result_data.get(sid_key)
                 if submission_id:
                     typer.echo(f"Got submission ID: {submission_id}")
-                    return self.get_result(submission_id)
+                    return self.get_test_result(submission_id)
                 else:
                     typer.echo("No submission ID in response", err=True)
                     return {"success": False, "error": "No submission ID received"}
@@ -177,7 +204,7 @@ def _format_output(self, output) -> str:
             return output.strip('[]"')
         return str(output)
 
-    def get_result(self, submission_id: str, timeout: int = 30) -> Dict[str, Any]:
+    def get_test_result(self, submission_id: str, timeout: int = 30) -> Dict[str, Any]:
         """Poll for results with timeout"""
         url = f"{self.BASE_URL}/submissions/detail/{submission_id}/check/"
         typer.echo(f"Polling for results at {url}")
@@ -216,7 +243,7 @@ def get_submission_result(self, submission_id: str) -> Dict[str, Any]:
         """Poll for submission results"""
         check_url = f"{self.BASE_URL}/submissions/detail/{submission_id}/check/"
 
-        for _ in range(20):  # Try for 20 seconds
+        for _ in range(20):  #
             response = self.session.get(check_url)
             if response.status_code != 200:
                 return {"success": False, "error": f"Failed to get results: {response.status_code}"}
@@ -234,50 +261,4 @@ def get_submission_result(self, submission_id: str) -> Dict[str, Any]:
 
             time.sleep(1)
 
-        return {"success": False, "error": "Timeout waiting for results"}
-
-    def get_test_result(self, interpret_id: str) -> Dict[str, Any]:
-        """Poll for test results"""
-        check_url = f"{self.BASE_URL}/submissions/detail/{interpret_id}/check/"
-
-        headers = {
-            'Accept': 'application/json',
-            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
-        }
-
-        max_attempts = 10
-        for attempt in range(max_attempts):
-            try:
-                response = self.session.get(check_url, headers=headers)
-                if response.status_code != 200:
-                    return {"success": False, "error": f"Failed to get results: {response.status_code}"}
-
-                result = response.json()
-
-                # Check for compilation errors first
-                if result.get('status_msg') == 'Compile Error':
-                    return {
-                        "success": False,
-                        "error": f"Compilation Error:\n{result.get('compile_error', 'Unknown compilation error')}"
-                    }
-
-                if result.get('state') == 'SUCCESS':
-                    status_msg = result.get('status_msg', 'Unknown')
-                    return {
-                        "success": True,
-                        "status": status_msg,
-                        "input": result.get('input', 'N/A'),
-                        "output": result.get('code_output', 'N/A').strip('[]"'),
-                        "expected": result.get('expected_output', 'N/A').strip('[]"'),
-                        "runtime": result.get('status_runtime', 'N/A'),
-                        "memory": result.get('memory', 'N/A'),
-                    }
-
-                # If not success yet, wait before next attempt
-                if attempt < max_attempts - 1:
-                    time.sleep(2)
-
-            except Exception as e:
-                return {"success": False, "error": f"Error checking results: {str(e)}"}
-
-        return {"success": False, "error": "Timeout waiting for test results"}
+        return {"success": False, "error": "Timeout waiting for results"}