#!/usr/bin/env python3
"""
mitmproxy addon: HTTP Basic Auth
Usage: mitmdump -s auth.py --set auth_user=admin --set auth_pass=secret
"""

from mitmproxy import http, ctx


class HTTPAuth:
    def __init__(self):
        self.auth_user = ctx.options.auth_user
        self.auth_pass = ctx.options.auth_pass
    
    def request(self, flow: http.HTTPFlow):
        # Skip auth for mitm.it (certificate download)
        if "mitm.it" in flow.request.pretty_host:
            return
        
        # Check for Authorization header
        auth_header = flow.request.headers.get("Authorization", "")
        
        if not auth_header.startswith("Basic "):
            self.challenge_auth(flow)
            return
        
        # Decode and verify credentials
        import base64
        try:
            credentials = base64.b64decode(auth_header[6:]).decode()
            username, password = credentials.split(":", 1)
            
            if username != self.auth_user or password != self.auth_pass:
                self.challenge_auth(flow)
        except Exception:
            self.challenge_auth(flow)
    
    def challenge_auth(self, flow: http.HTTPFlow):
        flow.response = http.Response.make(
            407,
            b"Proxy Authentication Required",
            {"Proxy-Authenticate": 'Basic realm="mitmproxy"'}
        )


addons = [HTTPAuth()]