From e5b7eea42f8436fb05521ccc57ddbc9ce4fbfd05 Mon Sep 17 00:00:00 2001
From: sultanofdisco <ohsujin712@gmail.com>
Date: Tue, 10 Jun 2025 01:37:11 +0900
Subject: [PATCH] =?UTF-8?q?nonceChecker=20=EC=88=98=EC=A0=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 addon/init.py        |  2 +-
 addon/nonce_check.py | 50 +++++++++++++++++++++++++++++++-------------
 2 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/addon/init.py b/addon/init.py
index c485248..f682683 100644
--- a/addon/init.py
+++ b/addon/init.py
@@ -57,7 +57,7 @@ class NonceAddon:
 
     async def response(self, flow: http.HTTPFlow):
         try:
-            await self.checker.response(flow)
+            await self.checker.check_nonce_in_id_token(flow)
         except Exception as e:
             print(f"[ERROR] NonceAddon failed: {e}")
             pass
diff --git a/addon/nonce_check.py b/addon/nonce_check.py
index bb1f379..a020ce4 100644
--- a/addon/nonce_check.py
+++ b/addon/nonce_check.py
@@ -8,20 +8,19 @@ from lib.report import save_report
 
 class NonceChecker:
     def is_oidc_flow(self, flow) -> bool:
-        req = flow.request
         res = flow.response
-        url = req.pretty_url
+        url = res.url
         parsed = urlparse(url)
-        query = parse_qs(parsed.query)
+        fragment_params = parse_qs(parsed.fragment)
         location = res.headers.get("location", "")
         content_type = res.headers.get("content-type", "")
 
-        if "/authorize" in url and "response_type" in query and "openid" in query.get("scope", [""])[0]:
-            return True
 
         if "application/json" in content_type:
             if "id_token" in res.text:
                 return True
+            if self.extract_id_token(self, flow):
+                return True
 
         if res.status_code in [302, 303]:
             if isinstance(location, list):
@@ -29,26 +28,29 @@ class NonceChecker:
             if "id_token=" in location:
                 return True
 
-        if "/authorize" in url and "nonce" in query:
+        if "id_token" in fragment_params:
             return True
 
         return False
 
 
-    def extract_id_token(self, response) -> Union[str, None]:
+    def extract_id_token(self, flow) -> Union[str, None]:
         """
         응답에서 id_token을 추출하는 함수.
         """
+        res = flow.response
         # 1. JSON 응답에 id_token 있음
         try:
-            if "application/json" in response.headers.get("content-type", ""):
-                data = response.json()
+            if "application/json" in res.headers.get("content-type", ""):
+                data = res.json()
                 return data.get("id_token")
+            else:
+                return None
         except Exception:
             pass
 
         # 2. Location 헤더에서 id_token 파싱 (예: #id_token=...&access_token=...)
-        location = response.headers.get("location", "")
+        location = res.headers.get("location", "")
         if location:
             if "#" in location:
                 fragment = location.split("#")[1]
@@ -62,22 +64,40 @@ class NonceChecker:
         return None
 
 
-    def decode_id_token(self, id_token: str) -> dict:
+    def decode_id_token(self, flow) -> dict:
+        res = flow.response
+        id_token = self.extract_id_token(res)
+        if not id_token:
+            return {}
         try:
             return jwt.decode(id_token, options={"verify_signature": False})
         except Exception as e:
             return {}
         
 
-    def check_nonce_in_id_token(self, flow, id_token: str) -> bool:
+    def check_nonce_in_id_token(self, flow) -> bool:
+        if not flow.response or not self.is_oidc_flow(flow):
+            # OIDC 플로우가 아니거나 응답이 없으면 nonce 체크를 건너뜀
+            return True
+        
+        res = flow.response
+        url = res.url
+        parsed = urlparse(url)
+        fragment_params = parse_qs(parsed.fragment)
+
+        if "id token" in fragment_params:
+            # id_token이 fragment에 있는 경우
+            id_token = fragment_params["id token"][0]
+            return True
+
+        id_token = self.extract_id_token(res)
         decoded = self.decode_id_token(id_token)
         nonce = decoded.get("nonce")
-        req = flow.request
-        url = req.pretty_url
+
         if not nonce:
             report_data = [{
                 'target': target.load(),
-                'status': "CRITICAL",
+                'status': "MEDIUM",
                 'title': "nonce is missing in id_token",
                 'description': "Nonce is present in the request but missing in the id_token.",
                 'uri': f"Original: {url}\nDecoded ID Token: {decoded}",