package main

import (
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/elazarl/goproxy"
	"golang.org/x/sys/windows/registry"
)

type Config struct {
	DomainsOfInterest []string `toml:"domains_of_interest"`
	Proxy             struct {
		Port int `toml:"port"`
	} `toml:"proxy"`
	Dump struct {
		OutputDir string `toml:"output_dir"`
		DOIDir    string `toml:"DOI_dir"`
	} `toml:"dump"`
}

type UserData struct {
	RequestBody  []byte
	ModifiedBody []byte
}

type ProxyServer struct {
	config        *Config
	tlsConfig     *tls.Config
	proxy         *goproxy.ProxyHttpServer
	server        *http.Server
	originalProxy string
	verbose       bool
	quiet         bool
}

func main() {
	// Parse command line flags
	var testConnectivity = flag.Bool("test", false, "Test proxy connectivity")
	var verbose = flag.Bool("v", false, "Enable verbose mode - dump all traffic instead of only modified requests/responses")
	var debugMode = flag.Bool("d", false, "Debug mode - dump modified requests/responses")
	flag.Parse()

	// Set console to UTF-8 on Windows to prevent garbled text
	if runtime.GOOS == "windows" {
		setConsoleUTF8()
	}

	fmt.Println("Reading dictionary...")
	dict, err := NewChineseDict("dict.txt")
	if err != nil {
		log.Fatalf("Failed to load dictionary: %v", err)
	}
	Dict = dict
	fmt.Printf("Dictionary loaded successfully, size=%d\n", Dict.GetCharacterCount())

	fmt.Println("Starting MITM proxy server...")

	// Load configuration
	config, err := loadConfig("config.toml")
	if err != nil {
		log.Fatalf("Failed to load configuration: %v", err)
	}
	log.Println(config.String())

	// Create output directories
	if err := os.MkdirAll(config.Dump.OutputDir, 0755); err != nil {
		log.Fatalf("Failed to create output directory: %v", err)
	}
	if config.Dump.DOIDir != "" {
		if err := os.MkdirAll(config.Dump.DOIDir, 0755); err != nil {
			log.Fatalf("Failed to create DOI directory: %v", err)
		}
	}

	// Create proxy server
	proxy, err := NewProxyServer(config, *verbose, *debugMode)
	if err != nil {
		log.Fatalf("Failed to create proxy server: %v", err)
	}

	// Install CA certificate
	if err := proxy.installCACert(); err != nil {
		log.Printf("Failed to install CA certificate: %v", err)
	} else {
		fmt.Println("CA certificate installed successfully")
	}

	// Set system proxy
	if err := proxy.setSystemProxy(); err != nil {
		log.Fatalf("Failed to set system proxy: %v", err)
	}
	fmt.Printf("System proxy set to 127.0.0.1:%d\n", config.Proxy.Port)

	// Verify proxy settings were applied
	if err := proxy.verifyProxySettings(); err != nil {
		fmt.Printf("Warning: Could not verify proxy settings: %v\n", err)
	} else {
		fmt.Printf("Proxy settings verified in Windows registry\n")
	}

	// Show current proxy configuration for debugging
	proxy.showProxyConfiguration()

	// Start proxy server
	serverStarted := make(chan error, 1)
	go func() {
		fmt.Printf("Starting proxy server on port %d...\n", config.Proxy.Port)
		if err := proxy.Start(); err != nil && err != http.ErrServerClosed {
			serverStarted <- err
			return
		}
		serverStarted <- nil
	}()

	// Wait for server to start or fail
	fmt.Printf("Waiting for server to start...\n")
	select {
	case err := <-serverStarted:
		if err != nil {
			log.Fatalf("Failed to start proxy server: %v", err)
		}
		fmt.Printf("Proxy server successfully started on port %d\n", config.Proxy.Port)
	case <-time.After(5 * time.Second):
		fmt.Printf("Proxy server appears to be starting (no immediate errors)\n")
	}

	// Run full connectivity test if requested
	if *testConnectivity {
		proxyAddr := fmt.Sprintf("127.0.0.1:%d", config.Proxy.Port)
		testProxyConnectivity(proxyAddr)
	}

	// Wait for interrupt signal
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	<-sigChan

	fmt.Println("\nShutting down proxy server...")
	proxy.Shutdown()
	fmt.Println("Proxy server closed")
}

func loadConfig(filename string) (*Config, error) {
	return parseConfig(filename)
}

func NewProxyServer(config *Config, verbose bool, debugMode bool) (*ProxyServer, error) {
	// Load hardcoded P12 certificate for MITM
	tlsConfig, err := loadHardcodedCertificate()
	if err != nil {
		return nil, fmt.Errorf("failed to load certificate: %v", err)
	}

	// Create goproxy instance
	goProxy := goproxy.NewProxyHttpServer()
	goProxy.Verbose = verbose

	ps := &ProxyServer{
		config:    config,
		tlsConfig: tlsConfig,
		proxy:     goProxy,
		verbose:   verbose,
		quiet:     !debugMode,
	}

	// Configure MITM for HTTPS traffic
	ps.setupMITM()

	// Setup request/response handlers
	ps.setupHandlers()

	// Create HTTP server
	ps.server = &http.Server{
		Addr:     fmt.Sprintf(":%d", config.Proxy.Port),
		Handler:  ps.proxy,
		ErrorLog: log.New(os.Stdout, "GOPROXY-SERVER: ", log.LstdFlags),
	}

	return ps, nil
}

// setupMITM configures HTTPS MITM using goproxy
func (p *ProxyServer) setupMITM() {
	// Set the CA certificate for MITM
	if len(p.tlsConfig.Certificates) > 0 {
		cert := p.tlsConfig.Certificates[0]
		goproxy.GoproxyCa = cert
		goproxy.OkConnect = &goproxy.ConnectAction{Action: goproxy.ConnectMitm, TLSConfig: goproxy.TLSConfigFromCA(&cert)}
	}

	// Enable MITM for all HTTPS connections
	p.proxy.OnRequest().HandleConnect(goproxy.AlwaysMitm)
}

// setupHandlers configures request/response handlers for logging and filtering
func (p *ProxyServer) setupHandlers() {
	// Add health check endpoint
	p.proxy.OnRequest().DoFunc(func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
		if r.Host == "127.0.0.1:"+fmt.Sprintf("%d", p.config.Proxy.Port) && r.URL.Path == "/proxy-health" {
			return r, goproxy.NewResponse(r, goproxy.ContentTypeText, http.StatusOK, "MITM Proxy is running")
		}
		return r, nil
	})

	// Log all HTTP requests and capture request body
	p.proxy.OnRequest().DoFunc(func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
		if !p.isDomainOfInterest(r.Host) && p.quiet {
			return r, nil
		}

		// Read request body once and recreate it for both dumping and forwarding
		if r.Body != nil {
			reqBody, err := io.ReadAll(r.Body)
			if err != nil {
				log.Printf("Failed to read request body: %v", err)
				return r, nil
			}
			r.Body.Close()

			var newReqBody []byte = nil

			if p.isDomainOfInterest(r.Host) {
				newReqBody, err = asrResultObfuscate(r, reqBody)
				if err != nil && err.Error() != "not an asr request" {
					log.Printf("Failed to obfuscate request body: %v", err)
				}
				if p.quiet && err == nil {
					log.Println("[INFO] ASR Request Body Modified")
				}
			}

			if newReqBody != nil {
				r.Body = io.NopCloser(bytes.NewReader(newReqBody))
				r.ContentLength = int64(len(newReqBody))
			} else {
				// Recreate the request body so it can be forwarded to the server
				r.Body = io.NopCloser(bytes.NewReader(reqBody))
				r.ContentLength = int64(len(reqBody))
			}

			// Store request body in context for later use in response handler
			if len(reqBody) > 0 {
				ctx.UserData = UserData{
					RequestBody:  reqBody,
					ModifiedBody: newReqBody,
				}
			}
		}

		return r, nil
	})

	// Log all HTTP responses and dump traffic
	p.proxy.OnResponse().DoFunc(func(r *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
		if p.quiet {
			return r
		}

		timestamp := time.Now().Format("20060102T15:04:05.000000")
		if r != nil {
			if p.verbose || p.isDomainOfInterest(ctx.Req.Host) {
				fmt.Printf(
					"[%s][INFO][Interest=%v] HTTP Response: %s %s\n",
					timestamp,
					p.isDomainOfInterest(ctx.Req.Host),
					r.Status,
					ctx.Req.URL.String(),
				)
			}

			// Get request body from context (if available)
			var reqBody []byte
			var modifiedBody []byte
			if ctx.UserData != nil {
				if userData, ok := ctx.UserData.(UserData); ok {
					reqBody = userData.RequestBody
					modifiedBody = userData.ModifiedBody
				}
			}

			// Read response body once and recreate it for both dumping and returning
			if r.Body != nil {
				respBody, err := io.ReadAll(r.Body)
				if err != nil {
					log.Printf("Failed to read response body: %v", err)
					return r
				}
				r.Body.Close()

				// Recreate the response body so it can be returned to client
				r.Body = io.NopCloser(bytes.NewReader(respBody))
				r.ContentLength = int64(len(respBody))

				// Dump traffic to file with both request and response bodies
				// Only dump if verbose mode is enabled OR if the request was modified
				if p.verbose || modifiedBody != nil {
					p.dumpHTTPTrafficWithBodies(ctx.Req, r, reqBody, modifiedBody, respBody)
				}
			} else {
				// No response body, but may have request body
				// Only dump if verbose mode is enabled OR if the request was modified
				if p.verbose || modifiedBody != nil {
					p.dumpHTTPTrafficWithBodies(ctx.Req, r, reqBody, modifiedBody, nil)
				}
			}
		}
		return r
	})
}

func (p *ProxyServer) Start() error {
	fmt.Printf("Starting HTTP server on %s\n", p.server.Addr)

	// Test if port is available
	listener, err := net.Listen("tcp", p.server.Addr)
	if err != nil {
		return fmt.Errorf("failed to bind to port %s: %v", p.server.Addr, err)
	}

	fmt.Printf("Successfully bound to port %s\n", p.server.Addr)

	// Start serving with our listener
	return p.server.Serve(listener)
}

func (p *ProxyServer) Shutdown() {
	// Restore system proxy settings
	if err := p.restoreSystemProxy(); err != nil {
		log.Printf("Failed to restore system proxy: %v", err)
	} else {
		fmt.Println("System proxy restored")
	}

	// Close HTTP server
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	p.server.Shutdown(ctx)
}

// dumpHTTPTrafficWithBodies dumps HTTP request and response with both bodies to file
func (p *ProxyServer) dumpHTTPTrafficWithBodies(req *http.Request, resp *http.Response, reqBody []byte, modifiedBody []byte, respBody []byte) {
	file, err := os.Create(p.getFilePath(req))
	if err != nil {
		log.Printf("Failed to create dump file: %v", err)
		return
	}
	defer file.Close()

	// Write request information
	fmt.Fprintf(file, "=== REQUEST ===\n")
	fmt.Fprintf(file, "%s %s %s\n", req.Method, req.URL.String(), req.Proto)
	fmt.Fprintf(file, "Host: %s\n", req.Host)

	// Write all request headers
	for name, values := range req.Header {
		for _, value := range values {
			fmt.Fprintf(file, "%s: %s\n", name, value)
		}
	}
	fmt.Fprintf(file, "\n")

	// Write request body
	if len(reqBody) > 0 {
		fmt.Fprintf(file, "%s\n", string(reqBody))
	}

	if modifiedBody != nil {
		fmt.Fprintf(file, "\n=== MODIFIED REQUEST BODY ===\n")
		fmt.Fprintf(file, "%s\n", string(modifiedBody))
	}

	// Write response information
	if resp != nil {
		fmt.Fprintf(file, "\n=== RESPONSE ===\n")
		fmt.Fprintf(file, "%s %s\n", resp.Proto, resp.Status)

		// Write all response headers
		for name, values := range resp.Header {
			for _, value := range values {
				fmt.Fprintf(file, "%s: %s\n", name, value)
			}
		}
		fmt.Fprintf(file, "\n")

		// Write response body
		if len(respBody) > 0 {
			fmt.Fprintf(file, "%s\n", string(respBody))
		}
	}
}

func (p *ProxyServer) getFilePath(req *http.Request) string {
	timestamp := time.Now().Format("20060102.150405.000000000")
	filename := fmt.Sprintf("%s_%s.txt", timestamp, sanitizeFilename(req.URL.String()))
	if p.isDomainOfInterest(req.Host) {
		return filepath.Join(p.config.Dump.DOIDir, filename)
	}
	return filepath.Join(p.config.Dump.OutputDir, filename)
}

// Keep the isDomainOfInterest method as it's still needed
func (p *ProxyServer) isDomainOfInterest(host string) bool {
	// Remove port number
	if colonIndex := strings.Index(host, ":"); colonIndex != -1 {
		host = host[:colonIndex]
	}

	for _, domain := range p.config.DomainsOfInterest {
		if strings.Contains(host, domain) {
			return true
		}
	}
	return false
}

// sanitizeFilename replaces unsafe characters for file names
func sanitizeFilename(filename string) string {
	// Replace unsafe characters
	replacer := strings.NewReplacer(
		":", "_",
		"/", "_",
		"\\", "_",
		"*", "_",
		"?", "_",
		"\"", "_",
		"<", "_",
		">", "_",
		"|", "_",
	)
	s := replacer.Replace(filename)
	if len(s) > 80 {
		return s[:80]
	}
	return s
}

func (p *ProxyServer) setSystemProxy() error {
	// Get current proxy settings
	k, err := registry.OpenKey(registry.CURRENT_USER, `Software\Microsoft\Windows\CurrentVersion\Internet Settings`, registry.QUERY_VALUE)
	if err != nil {
		return err
	}
	defer k.Close()

	// Read current ProxyServer settings
	currentProxy, _, err := k.GetStringValue("ProxyServer")
	if err == nil {
		p.originalProxy = currentProxy
	}

	// Set new proxy
	k, err = registry.OpenKey(registry.CURRENT_USER, `Software\Microsoft\Windows\CurrentVersion\Internet Settings`, registry.SET_VALUE)
	if err != nil {
		return err
	}
	defer k.Close()

	proxyAddr := fmt.Sprintf("127.0.0.1:%d", p.config.Proxy.Port)

	// Enable proxy
	if err := k.SetDWordValue("ProxyEnable", 1); err != nil {
		return err
	}

	// Set proxy server address for BOTH HTTP and HTTPS
	// Format: "http=proxy:port;https=proxy:port" or just "proxy:port" for both
	if err := k.SetStringValue("ProxyServer", proxyAddr); err != nil {
		return err
	}

	// Ensure no proxy bypass for local addresses
	if err := k.SetStringValue("ProxyOverride", "<local>"); err != nil {
		// This is not critical, continue anyway
		fmt.Printf("Warning: Could not set ProxyOverride: %v\n", err)
	}

	// Refresh IE settings and notify system
	return p.refreshIESettings()
}

func (p *ProxyServer) restoreSystemProxy() error {
	k, err := registry.OpenKey(registry.CURRENT_USER, `Software\Microsoft\Windows\CurrentVersion\Internet Settings`, registry.SET_VALUE)
	if err != nil {
		return err
	}
	defer k.Close()

	// Disable proxy
	if err := k.SetDWordValue("ProxyEnable", 0); err != nil {
		return err
	}

	// Restore original proxy settings
	if p.originalProxy != "" {
		if err := k.SetStringValue("ProxyServer", p.originalProxy); err != nil {
			return err
		}
	}

	// Refresh IE settings
	return p.refreshIESettings()
}

func (p *ProxyServer) refreshIESettings() error {
	// Multiple methods to ensure Windows recognizes the proxy change

	// Method 1: Reset WinHTTP proxy
	cmd1 := exec.Command("cmd", "/c", "netsh winhttp reset proxy")
	if err := cmd1.Run(); err != nil {
		fmt.Printf("Warning: netsh winhttp reset failed: %v\n", err)
	}

	// Method 2: Import IE proxy settings to WinHTTP
	cmd2 := exec.Command("cmd", "/c", "netsh winhttp import proxy source=ie")
	if err := cmd2.Run(); err != nil {
		fmt.Printf("Warning: netsh winhttp import failed: %v\n", err)
	}

	// Method 3: Force refresh Internet Options
	cmd3 := exec.Command("cmd", "/c", "rundll32.exe inetcpl.cpl,ClearMyTracksByProcess 8")
	if err := cmd3.Run(); err != nil {
		fmt.Printf("Warning: Internet Options refresh failed: %v\n", err)
	}

	return nil
}

func (p *ProxyServer) verifyProxySettings() error {
	k, err := registry.OpenKey(registry.CURRENT_USER, `Software\Microsoft\Windows\CurrentVersion\Internet Settings`, registry.QUERY_VALUE)
	if err != nil {
		return fmt.Errorf("failed to open registry key: %v", err)
	}
	defer k.Close()

	// Check if proxy is enabled
	proxyEnable, _, err := k.GetIntegerValue("ProxyEnable")
	if err != nil {
		return fmt.Errorf("failed to read ProxyEnable: %v", err)
	}
	if proxyEnable != 1 {
		return fmt.Errorf("proxy is not enabled (ProxyEnable = %d)", proxyEnable)
	}

	// Check proxy server setting
	proxyServer, _, err := k.GetStringValue("ProxyServer")
	if err != nil {
		return fmt.Errorf("failed to read ProxyServer: %v", err)
	}

	expectedProxy := fmt.Sprintf("127.0.0.1:%d", p.config.Proxy.Port)
	if proxyServer != expectedProxy {
		return fmt.Errorf("proxy server mismatch: expected %s, got %s", expectedProxy, proxyServer)
	}

	return nil
}

func (p *ProxyServer) showProxyConfiguration() {
	fmt.Printf("Current Proxy Configuration:\n")

	// Show Windows proxy settings
	k, err := registry.OpenKey(registry.CURRENT_USER, `Software\Microsoft\Windows\CurrentVersion\Internet Settings`, registry.QUERY_VALUE)
	if err != nil {
		fmt.Printf("Could not read proxy settings: %v\n", err)
		return
	}
	defer k.Close()

	if proxyEnable, _, err := k.GetIntegerValue("ProxyEnable"); err == nil {
		fmt.Printf("   ProxyEnable: %d\n", proxyEnable)
	}

	if proxyServer, _, err := k.GetStringValue("ProxyServer"); err == nil {
		fmt.Printf("   ProxyServer: %s\n", proxyServer)
	}

	if proxyOverride, _, err := k.GetStringValue("ProxyOverride"); err == nil {
		fmt.Printf("   ProxyOverride: %s\n", proxyOverride)
	}

	// Show WinHTTP proxy settings
	fmt.Printf("   WinHTTP Proxy Settings:\n")
	cmd := exec.Command("cmd", "/c", "netsh winhttp show proxy")
	if output, err := cmd.CombinedOutput(); err == nil {
		fmt.Printf("   %s\n", string(output))
	} else {
		fmt.Printf("   ❌ Could not get WinHTTP proxy settings: %v\n", err)
	}
}

func (p *ProxyServer) installCACert() error {
	// Use hardcoded CA certificate data
	certData := []byte(getHardcodedCACert())

	// Parse certificate
	var cert *x509.Certificate
	var err error

	// Try parsing DER format directly
	cert, err = x509.ParseCertificate(certData)
	if err != nil {
		// Try parsing PEM format
		block, _ := pem.Decode(certData)
		if block == nil {
			return fmt.Errorf("unable to parse certificate file")
		}
		cert, err = x509.ParseCertificate(block.Bytes)
		if err != nil {
			return err
		}
	}

	// Write certificate to temporary file
	tempFile, err := os.CreateTemp("", "ca_cert_*.crt")
	if err != nil {
		return err
	}
	tempFileName := tempFile.Name()
	defer os.Remove(tempFileName)

	if err := pem.Encode(tempFile, &pem.Block{
		Type:  "CERTIFICATE",
		Bytes: cert.Raw,
	}); err != nil {
		tempFile.Close()
		return err
	}

	// Close the file before certutil accesses it to avoid sharing violation
	tempFile.Close()

	// Use certutil to install certificate to trusted root certificate authorities
	cmd := exec.Command("certutil", "-addstore", "-f", "Root", tempFileName)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed to install certificate: %v, output: %s", err, string(output))
	}

	return nil
}

// setConsoleUTF8 sets the Windows console to UTF-8 code page to prevent garbled text
func setConsoleUTF8() {
	// Set console input and output code page to UTF-8 (65001)
	cmd := exec.Command("cmd", "/c", "chcp 65001 >nul")
	cmd.Run()
}

// testProxyConnectivity tests the proxy with simple HTTP and HTTPS requests
func testProxyConnectivity(proxyAddr string) {
	fmt.Println("Testing proxy connectivity...")

	// Create proxy URL
	proxyURL, err := url.Parse(fmt.Sprintf("http://%s", proxyAddr))
	if err != nil {
		fmt.Printf("Failed to parse proxy URL: %v\n", err)
		return
	}

	// Create HTTP client with proxy
	transport := &http.Transport{
		Proxy: http.ProxyURL(proxyURL),
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
	}
	client := &http.Client{
		Transport: transport,
		Timeout:   30 * time.Second,
	}

	// Test HTTP request
	fmt.Println("Testing HTTP request...")
	testHTTPRequest(client, "http://httpbin.org/ip")

	// Test HTTPS request
	fmt.Println("Testing HTTPS request...")
	testHTTPSRequest(client, "https://httpbin.org/ip")

	// Test redirects (should be transparent)
	fmt.Println("Testing HTTP redirect...")
	testRedirectTransparency(client, "http://httpbin.org/redirect/1")

	// Test domain that might cause issues
	fmt.Println("Testing Google...")
	testHTTPSRequest(client, "https://www.google.com")
}

func testHTTPRequest(client *http.Client, testURL string) {
	resp, err := client.Get(testURL)
	if err != nil {
		fmt.Printf("HTTP test failed: %v\n", err)
		return
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		fmt.Printf("Failed to read HTTP response: %v\n", err)
		return
	}

	fmt.Printf("HTTP test successful: %s\n", resp.Status)
	fmt.Printf("   Response: %s\n", string(body)[:min(len(body), 200)])
}

func testHTTPSRequest(client *http.Client, testURL string) {
	resp, err := client.Get(testURL)
	if err != nil {
		fmt.Printf("HTTPS test failed: %v\n", err)
		return
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		fmt.Printf("Failed to read HTTPS response: %v\n", err)
		return
	}

	fmt.Printf("HTTPS test successful: %s\n", resp.Status)
	fmt.Printf("   Response: %s\n", string(body)[:min(len(body), 200)])
}

func testRedirectTransparency(client *http.Client, testURL string) {
	// Create a client that doesn't follow redirects to test transparency
	testClient := &http.Client{
		Transport: client.Transport,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
		Timeout: client.Timeout,
	}

	resp, err := testClient.Get(testURL)
	if err != nil {
		fmt.Printf("Redirect transparency test failed: %v\n", err)
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 300 && resp.StatusCode < 400 {
		location := resp.Header.Get("Location")
		fmt.Printf("Redirect transparency test successful: %s -> %s\n", resp.Status, location)
	} else {
		fmt.Printf("Expected redirect, got: %s (this may be normal)\n", resp.Status)
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}