// grpc-mitm.go
package main

import (
	"context"
	"crypto/tls"
	"encoding/base64"
	"encoding/hex"
	"flag"
	"fmt"
	// "io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"strings"
	"time"

	"golang.org/x/net/http2"
	"golang.org/x/net/http2/h2c"
)

var (
	listenAddr = flag.String("listen", ":9191", "listen address (h2c)")
	backend    = flag.String("backend", "127.0.0.1:9095", "backend gRPC (host:port)")
	maxBody    = flag.Int("max-body", 10*1024*1024, "max body (bytes) to buffer/log")
	
	// Match & Replace rules for checksums (base64 strings)
	replaceRules = map[string]string{
		// "ewkem4QbzGXNj3bQTWVIiQ==": "wvmtNtAAFE3RME6L+qVWEg==", // OK Pour DATA.7.0.3.x64
		// administrator 7.0.3 legit : QWfKsZTOryjsBqqBywhuXQ==
		// "NWwiznM/6zaFLh+wzLiaUQ==":"zVPZfzltIkT/fFJAdrfTHQ==", // adm 7.0.0. x64
		// "M8HdeQiMcukzSkF+aevHvQ==":"mLfRaen7D+fcEPls+gT0sQ==", // adm 5.1.6. anycpu
		"nsPH/kDofBLRHtiT9gf7Hw==":"gIYMdbsrrZZx+ntRYhDu2g==", // adm 6.0.0. anycpu
	}
	
	// Replace file sizes (protobuf field 11, varint encoded)
	// Format: old varint bytes -> new varint bytes
	replaceSizes = map[string]string{
		// Data 7.0.0: 148836023 bytes -> Data 7.0.3: 149094004 bytes
		// 148836023 = 0x8df0eb7 -> varint: b7 9d fc 46
		// 149094004 = 0x8e2fe74 -> varint: f4 fc 8b 47
		
		// "\xf4\xfc\x8b\x47": "\xe8\xfd\x8b\x47", // OK Pour DATA.7.0.3.x64
		// "\xe7\xc5\xdc\x86\x01":"\xd0\xf2\xdc\x86\x01", // adm 7.0.0. x64
		// "\xb7\xcc\x87\x27":"\xa0\xf9\x87\x27", // adm 5.1.6
		"\xb5\xf6\x9a\x2b":"\x9e\xa3\x9b\x2b", // adm 6.0.0
	}
)

func main() {
	flag.Parse()
	
	// Detect backend protocol and strip prefix
	backendHost := *backend
	useTLS := false
	
	// gRPC typically uses TLS unless explicitly http:// or grpc-insecure://
	if strings.HasPrefix(backendHost, "grpc://") {
		useTLS = true
		backendHost = strings.TrimPrefix(backendHost, "grpc://")
	} else if strings.HasPrefix(backendHost, "https://") {
		useTLS = true
		backendHost = strings.TrimPrefix(backendHost, "https://")
	} else if strings.HasPrefix(backendHost, "http://") {
		useTLS = false
		backendHost = strings.TrimPrefix(backendHost, "http://")
	} else if strings.HasPrefix(backendHost, "grpc-insecure://") {
		useTLS = false
		backendHost = strings.TrimPrefix(backendHost, "grpc-insecure://")
	} else {
		// Default: if port 443 or standard gRPC ports, assume TLS
		useTLS = strings.Contains(backendHost, ":443") || strings.Contains(backendHost, ":9095")
	}
	
	scheme := "http"
	if useTLS {
		scheme = "https"
	}
	
	target := scheme + "://" + backendHost
	
	// Extract hostname without port for SNI
	hostname := backendHost
	if idx := strings.LastIndex(hostname, ":"); idx != -1 {
		hostname = hostname[:idx]
	}
	
	log.Printf("[mitm] Backend config: host=%s, hostname=%s, TLS=%v, scheme=%s", backendHost, hostname, useTLS, scheme)

	// Reverse proxy using standard ReverseProxy but with custom transport for http2
	proxy := &httputil.ReverseProxy{
		Director: func(req *http.Request) {
			// rewrite URL to target backend
			req.URL.Scheme = scheme
			req.URL.Host = backendHost
			// leave :authority/Host as-is (we want to forward Host maybe)
		},
		ModifyResponse: func(resp *http.Response) error {
			log.Printf("[mitm] Got response: status=%d", resp.StatusCode)
			// read and log response body
			if resp.Body != nil {
				buf, err := ioutil.ReadAll(resp.Body)
				if err != nil {
					log.Printf("[mitm] error reading response body: %v", err)
					resp.Body = ioutil.NopCloser(strings.NewReader(""))
					return nil
				}
				
				// Close original body
				resp.Body.Close()
				
				// Apply replacement rules on the binary data
				originalBuf := buf
				modified := false
				
				// 1. Replace file sizes FIRST (avant de modifier les checksums)
				//    Car les checksums sont des strings qui peuvent contenir n'importe quels bytes
				for oldBytes, newBytes := range replaceSizes {
					oldBytesSlice := []byte(oldBytes)
					newBytesSlice := []byte(newBytes)
					
					log.Printf("[mitm] Searching for size pattern: %x", oldBytesSlice)
					
					if idx := findBytes(buf, oldBytesSlice); idx != -1 {
						log.Printf("[mitm] Found size pattern at offset %d", idx)
						
						// Vérifier que c'est précédé de 0x58 (field 11)
						if idx > 0 && buf[idx-1] == 0x58 {
							log.Printf("[mitm] Confirmed: preceded by 0x58 (field 11)")
							
							// Gérer les longueurs différentes
							if len(oldBytesSlice) == len(newBytesSlice) {
								// Même longueur: remplacement simple
								log.Printf("[mitm] Replacing %x with %x (same length)", oldBytesSlice, newBytesSlice)
								copy(buf[idx:], newBytesSlice)
							} else {
								// Longueur différente: reconstruire le buffer
								newBuf := make([]byte, 0, len(buf)-len(oldBytesSlice)+len(newBytesSlice))
								newBuf = append(newBuf, buf[:idx]...)
								newBuf = append(newBuf, newBytesSlice...)
								newBuf = append(newBuf, buf[idx+len(oldBytesSlice):]...)
								buf = newBuf
								log.Printf("[mitm] ✓ REPLACED file size (length changed: %d → %d bytes)", 
									len(oldBytesSlice), len(newBytesSlice))
							}
							modified = true
							log.Printf("[mitm] ✓ REPLACED file size at offset %d", idx)
						} else {
							log.Printf("[mitm] Pattern found but NOT preceded by 0x58 (found: 0x%x)", buf[idx-1])
						}
					} else {
						log.Printf("[mitm] Size pattern NOT found in buffer")
					}
				}
				
				// 2. Replace checksums AFTER (sur le buffer potentiellement modifié)
				for oldStr, newStr := range replaceRules {
					oldBytes := []byte(oldStr)
					newBytes := []byte(newStr)
					if idx := findBytes(buf, oldBytes); idx != -1 {
						if len(oldBytes) == len(newBytes) {
							copy(buf[idx:], newBytes)
						} else {
							newBuf := make([]byte, 0, len(buf)-len(oldBytes)+len(newBytes))
							newBuf = append(newBuf, buf[:idx]...)
							newBuf = append(newBuf, newBytes...)
							newBuf = append(newBuf, buf[idx+len(oldBytes):]...)
							buf = newBuf
						}
						modified = true
						log.Printf("[mitm] ✓ REPLACED checksum: %q → %q (at offset %d)", 
							oldStr, newStr, idx)
					}
				}
				
				if modified {
					log.Printf("[mitm] Response modified: %d bytes → %d bytes", len(originalBuf), len(buf))
				}
				
				// Restore full body for client (modified or original)
				resp.Body = ioutil.NopCloser(strings.NewReader(string(buf)))
				
				// Log only up to maxBody
				logSize := len(buf)
				if logSize > *maxBody {
					logSize = *maxBody
				}
				
				if len(buf) > 0 {
					log.Printf("---- RESPONSE for %s (%d bytes) ----", resp.Request.URL.Path, len(buf))
					logHeaders(resp.Header)
					logPayload(buf[:logSize])
					log.Printf("---- END RESPONSE ----")
				} else {
					log.Printf("[mitm] response for %s (empty body)", resp.Request.URL.Path)
				}
			}
			return nil
		},
		Transport: &http2TransportWrapper{
			inner: &http2.Transport{
				AllowHTTP: !useTLS, // allow plain h2 only if not using TLS
				DialTLSContext: func(ctx context.Context, network, addr string, cfg *tls.Config) (net.Conn, error) {
					log.Printf("[mitm] Dialing backend: %s %s (TLS=%v)", network, addr, useTLS)
					if useTLS {
						// For gRPC with TLS, use actual TLS connection
						if cfg == nil {
							cfg = &tls.Config{}
						}
						// Ensure SNI is set correctly
						if cfg.ServerName == "" {
							cfg.ServerName = hostname
						}
						cfg.InsecureSkipVerify = true // For testing - set to false in production
						cfg.NextProtos = []string{"h2"} // ALPN for HTTP/2 (required for gRPC)
						
						log.Printf("[mitm] Attempting TLS handshake with SNI=%s, ALPN=%v", cfg.ServerName, cfg.NextProtos)
						
						dialer := &net.Dialer{
							Timeout:   10 * time.Second,
							KeepAlive: 30 * time.Second,
						}
						
						conn, err := tls.DialWithDialer(dialer, network, addr, cfg)
						if err != nil {
							log.Printf("[mitm] TLS dial error: %v", err)
							return nil, err
						}
						
						// Verify connection state
						state := conn.ConnectionState()
						log.Printf("[mitm] TLS connection established to %s (negotiated: %s, version: %x)", 
							addr, state.NegotiatedProtocol, state.Version)
						return conn, nil
					}
					// For h2c (gRPC without TLS), dial plaintext TCP
					conn, err := net.DialTimeout(network, addr, 10*time.Second)
					if err != nil {
						log.Printf("[mitm] TCP dial error: %v", err)
						return nil, err
					}
					log.Printf("[mitm] TCP connection established to %s", addr)
					return conn, nil
				},
				TLSClientConfig: &tls.Config{
					InsecureSkipVerify: true,
					ServerName:         hostname,
					NextProtos:         []string{"h2"}, // ALPN for HTTP/2
				},
			},
			scheme: scheme,
		},
		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
			log.Printf("[mitm] PROXY ERROR: %v", err)
			http.Error(w, fmt.Sprintf("Proxy error: %v", err), http.StatusBadGateway)
		},
	}

	// server that accepts h2c (HTTP/2 without TLS)
	handler := h2c.NewHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// log incoming request headers + metadata
		log.Printf(">>> REQUEST %s %s from %s", r.Proto, r.URL.Path, r.RemoteAddr)
		logHeaders(r.Header)

		// read body and log it (works well for unary calls)
		if r.Body != nil {
			bodyBuf, err := ioutil.ReadAll(r.Body)
			if err != nil {
				log.Printf("[mitm] error reading request body: %v", err)
				http.Error(w, "cannot read body", http.StatusBadRequest)
				return
			}
			
			// Close original body
			r.Body.Close()
			
			// Restore body for proxying
			r.Body = ioutil.NopCloser(strings.NewReader(string(bodyBuf)))

			if len(bodyBuf) > 0 {
				log.Printf("---- REQUEST PAYLOAD (%d bytes) ----", len(bodyBuf))
				logPayload(bodyBuf)
				log.Printf("---- END REQUEST PAYLOAD ----")
			} else {
				log.Printf("[mitm] request payload empty")
			}
		}

		// forward to backend
		proxy.ServeHTTP(w, r)
	}), &http2.Server{})

	srv := &http.Server{
		Addr:         *listenAddr,
		Handler:      handler,
		ReadTimeout:  0,
		WriteTimeout: 0,
	}

	log.Printf("gRPC MITM listening %s -> backend %s (max-body=%d bytes)", *listenAddr, target, *maxBody)
	if err := srv.ListenAndServe(); err != nil {
		log.Fatalf("server error: %v", err)
	}
}

// http2TransportWrapper implements RoundTripper by dialing http (plaintext) over TCP
type http2TransportWrapper struct {
	inner  *http2.Transport
	scheme string
}

func (t *http2TransportWrapper) RoundTrip(req *http.Request) (*http.Response, error) {
	// Clear RequestURI - it's set by the proxy but must be empty for client requests
	req.RequestURI = ""
	
	// Ensure proper URL
	req.URL.Scheme = t.scheme
	if req.URL.Host == "" {
		req.URL.Host = req.Host
	}
	
	log.Printf("[mitm] RoundTrip: %s %s", req.Method, req.URL.String())
	
	// Use inner transport's RoundTrip directly
	resp, err := t.inner.RoundTrip(req)
	if err != nil {
		log.Printf("[mitm] RoundTrip error: %v", err)
		return nil, err
	}
	log.Printf("[mitm] RoundTrip success: status=%d", resp.StatusCode)
	return resp, nil
}

func logHeaders(h http.Header) {
	for k, v := range h {
		// mask sensitive headers if needed
		log.Printf("  %s: %s", k, strings.Join(v, ", "))
	}
}

func logPayload(b []byte) {
	// binary: show base64 and hex and a printable prefix (if ascii)
	b64 := base64.StdEncoding.EncodeToString(b)
	hexs := hex.EncodeToString(b)
	printable := toPrintable(b, 200)
	log.Printf("  base64(len=%d): %s", len(b), b64)
	log.Printf("  hex(len=%d): %s", len(b), hexs)
	if len(printable) > 0 {
		log.Printf("  ascii prefix: %s", printable)
	}
}

func toPrintable(b []byte, max int) string {
	n := len(b)
	if n > max {
		n = max
	}
	out := make([]byte, 0, n)
	for i := 0; i < n; i++ {
		c := b[i]
		if c >= 32 && c <= 126 {
			out = append(out, c)
		} else {
			out = append(out, '.')
		}
	}
	return string(out)
}

// findBytes searches for needle in haystack and returns the index, or -1 if not found
func findBytes(haystack, needle []byte) int {
	if len(needle) == 0 {
		return 0
	}
	if len(needle) > len(haystack) {
		return -1
	}
	for i := 0; i <= len(haystack)-len(needle); i++ {
		found := true
		for j := 0; j < len(needle); j++ {
			if haystack[i+j] != needle[j] {
				found = false
				break
			}
		}
		if found {
			return i
		}
	}
	return -1
}