import jwt
from urllib.parse import urlparse, parse_qs
from typing import Union

from lib.report_vuln import report_vuln

class NonceChecker:
    def is_oidc_flow(self, flow) -> bool:
        req = flow.request
        res = flow.response
        url = req.pretty_url
        parsed = urlparse(url)
        query = parse_qs(parsed.query)
        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 res.status_code in [302, 303]:
            if isinstance(location, list):
                location = location[0]
            if "id_token=" in location:
                return True

        if "/authorize" in url and "nonce" in query:
            return True

        return False


    def extract_id_token(self, response) -> Union[str, None]:
        """
        응답에서 id_token을 추출하는 함수.
        """
        # 1. JSON 응답에 id_token 있음
        try:
            if "application/json" in response.headers.get("content-type", ""):
                data = response.json()
                return data.get("id_token")
        except Exception:
            pass

        # 2. Location 헤더에서 id_token 파싱 (예: #id_token=...&access_token=...)
        location = response.headers.get("location", "")
        if location:
            if "#" in location:
                fragment = location.split("#")[1]
                params = parse_qs(fragment)
                return params.get("id_token", [None])[0]
            elif "?" in location:
                query = location.split("?")[1]
                params = parse_qs(query)
                return params.get("id_token", [None])[0]

        return None


    def decode_id_token(self, id_token: str) -> dict:
        try:
            return jwt.decode(id_token, options={"verify_signature": False})
        except Exception as e:
            return {}
        
    # TODO id_token을 파싱하는 부분이 누락되어있습니다.
    def check_nonce_in_id_token(self, flow, id_token: str) -> bool:
        decoded = self.decode_id_token(id_token)
        nonce = decoded.get("nonce")
        if not nonce:
            report_vuln(
                title="Nonce Check Failed",
                desc="id_token에 nonce가 없습니다.",
                status="HIGH",
                uri=flow.request.url
            )
            return False
        else:
            return True