[Refactor and Enhance] addon init.py의 비동기 작업을 더욱 효율적으로 수행

addon/access_token.py

@@ -1,6 +1,5 @@
 import re
-from dataclasses import dataclass, asdict
-from typing import List, Dict, Optional, Any
+from typing import Optional, Any
 import asyncio
 
 from mitmproxy.http import HTTPFlow

addon/csrf_check.py

@@ -1,5 +1,5 @@
 # csrf_check.py
-from mitmproxy import http, ctx
+from mitmproxy import http
 from urllib.parse import urlparse, parse_qs, unquote
 import httpx
 from typing import Optional, Union, List

addon/init.py

@@ -9,6 +9,7 @@ from access_token import AccessTokenScanner
 from addon.google_login_hint import GoogleLoginHint
 import os
 from dotenv import load_dotenv
+from lib.utils.try_catch import try_catch
 from lib.false_true_varifing_task import FalseTrueVarifingTask
 
 # Initialize the singleton task manager
@@ -16,111 +17,39 @@ false_true_varifing_task = FalseTrueVarifingTask()
 
 load_dotenv(override=True)
 
-class PKCEAddon:
-    def __init__(self):
-        self.checker = PKCEDowngradeChecker()
-
-    async def request(self, flow: http.HTTPFlow):
-        print(
-            f"[DEBUG] Processing request: {flow.request.method} {flow.request.pretty_url}"
-        )
-        try:
-            # 오탐 검사하고 있을때는 검증하지 않음
-            if false_true_varifing_task.is_verifing_false_true():
-                return
-                     
-            await self.checker.test(flow)
-        except Exception as e:
-            print(f"[ERROR] Addon failed: {e}")
-            pass
-            
-            
-class CsrfAddon:
-    def __init__(self):
-        self.checker = CsrfChecker()
-
-    async def response(self, flow: http.HTTPFlow):
-        try:
-            # 오탐 검사하고 있을때는 검증하지 않음
-            if false_true_varifing_task.is_verifing_false_true():
-                return
-            await self.checker.response(flow)
-        except Exception as e:
-            print(f"[ERROR] CSRF Addon failed: {e}")
-            pass
-        
-        
-class ScopeAddon:
-    def __init__(self):
-        self.checker = ScopeDetection()
-
-    async def response(self, flow: http.HTTPFlow):
-        try:
-            # 오탐 검사하고 있을때는 검증하지 않음
-            if false_true_varifing_task.is_verifing_false_true():
-                return
-            await self.checker.test(flow)
-        except Exception as e:
-            print(f"[ERROR] ScopeDetection failed: {e}")
-
-
-
-class NonceAddon:
-    def __init__(self):
-        self.checker = NonceChecker()
-
-    async def response(self, flow: http.HTTPFlow):
-        try:
-            pass
-            # TODO id_token을 파싱하는 부분이 누락되어있습니다.
-            # await self.checker.check_nonce_in_id_token(flow)
-        except Exception as e:
-            print(f"[ERROR] NonceAddon failed: {e}")
-            pass
-
-
-
-class AccessTokenAddon:
-    def __init__(self):
-        self.checker = AccessTokenScanner()
-
-    async def response(self, flow: http.HTTPFlow):
-        try:
-            # 오탐 검사하고 있을때는 검증하지 않음
-            if false_true_varifing_task.is_verifing_false_true():
-                return
-            await self.checker.scan(flow)
-        except Exception as e:
-            print(f"[ERROR] AccessToken Addon failed: {e}")
-            pass
-
-class RedirectBypassAddon:
-    def __init__(self):
-        self.checker = RedirectBypassChecker()
-
-    # request 대신 response 로 바꿔 보세요:
-    async def response(self, flow: http.HTTPFlow):
-        try:
-            # 오탐 검사하고 있을때는 검증하지 않음
-            if false_true_varifing_task.is_verifing_false_true():
-                return
-            await self.checker.test(flow)
-        except Exception as e:
-            print(f"[ERROR] RedirectBypass Addon failed: {e}")
-
-class GoogleLoginHintAddon():
+class AddonBase:
+    """
+    Base class for addons.
+    Each addon should implement its own request or response method.
+    """
+    
     def __init__(self) -> None:
         if os.getenv('GOOGLE_ID'):
-            self.checker = GoogleLoginHint()
+            self.google_login_hint = GoogleLoginHint()
         else:
-            self.checker = None
-    
-    async def request(self, flow: http.HTTPFlow):
-        if self.checker is None:
-            return
-        try:
-            await self.checker.request(flow)
-        except Exception as e:
-            print(f"[ERROR] GoogleLoginHint Addon failed: {e}")
+            self.google_login_hint = None
 
-addons = [PKCEAddon(), ScopeAddon(), CsrfAddon(), NonceAddon(), AccessTokenAddon(), GoogleLoginHintAddon(), RedirectBypassAddon()]
+    async def request(self, flow: http.HTTPFlow):
+        if false_true_varifing_task.is_verifing_false_true():
+            return
+            
+        tasks = [
+            try_catch(self.google_login_hint.request(flow)) if self.google_login_hint else None,
+            try_catch(PKCEDowngradeChecker().test(flow)),
+        ]
+        await asyncio.gather(*tasks)
+
+    async def response(self, flow: http.HTTPFlow):
+        if false_true_varifing_task.is_verifing_false_true():
+            return
+            
+        tasks = [
+            try_catch(CsrfChecker().response(flow)),
+            try_catch(ScopeDetection().test(flow)),
+            # try_catch(NonceChecker().check_nonce_in_request(flow)),
+            try_catch(AccessTokenScanner().scan(flow)),
+            try_catch(RedirectBypassChecker().test(flow)),
+        ]
+        await asyncio.gather(*tasks)
+
+addons = [AddonBase()]

addon/pkce_check.py

@@ -1,9 +1,6 @@
 from urllib.parse import urlparse, parse_qs, urlencode, urlunparse
-import asyncio
 import httpx
 from typing import Dict, List
-
-import lib.cur_target_url as cur_target_url
 from lib.report_vuln import report_vuln
 
 

addon/scope_detection.py

@@ -12,18 +12,16 @@ class ScopeDetection:
             return scope_values[0]
         return None
 
-    async def check_scope(self, flow):
+    async def test(self, flow):
+        if not is_oauth_uri(flow.request.pretty_url):
+            return
+        
         req = flow.request
-        res = flow.response
 
         parsed = urlparse(req.pretty_url)
         query = parsed.query
 
-        location = res.headers.get("Location", "")
-        location_query = urlparse(location).query
-
         query_scope = self.get_scope_from_query(query)
-        location_scope = self.get_scope_from_query(location_query)
 
         if query_scope in ["all", "*"]:
             report_vuln(
@@ -32,17 +30,3 @@ class ScopeDetection:
                 status="WARNING",
                 uri=req.pretty_url
             )
-        if location_scope in ["all", "*"]:
-            report_vuln(
-                title="OAuth Scope Value Issue",
-                desc=f"Scope value issue detected in response location: {location_scope}",
-                status="WARNING",
-                uri=location
-            )
-
-    async def test(self, flow):
-        
-        if not is_oauth_uri(flow.request.pretty_url):
-            return
-
-        await self.check_scope(flow)