package main

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/big"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"
)

// Config
var (
	TargetHTTPHost = "https://api.asmodee.net"
	TargetTCPHost  = "got.games.asmodee.net"
	TargetTCPPort  = "2445"
	ProxyHTTPPort  = ":8080"
	ProxyTCPPort   = ":3000"
	LogFileName    = ""
	LogFile        *os.File
	LogMutex       sync.Mutex
)

// Log Entry Structure
type LogEntry struct {
	Timestamp string      `json:"timestamp"`
	Protocol  string      `json:"protocol"` // HTTP or TCP
	Type      string      `json:"type"`     // REQUEST or RESPONSE
	Summary   string      `json:"summary"`  // e.g. "GET /v1/foo" or "Packet 123"
	Data      interface{} `json:"data"`     // Structured data or raw info
}

// Write to JSONL
func logJSON(entry LogEntry) {
	entry.Timestamp = time.Now().Format(time.RFC3339Nano)
	LogMutex.Lock()
	defer LogMutex.Unlock()

	bytes, err := json.Marshal(entry)
	if err == nil {
		if LogFile != nil {
			LogFile.Write(bytes)
			LogFile.WriteString("\n")
		} else {
			fmt.Println(string(bytes))
		}
	}
}

func main() {
	flag.StringVar(&TargetHTTPHost, "target-http", "https://api.asmodee.net", "Target HTTP Host")
	flag.StringVar(&TargetTCPHost, "target-tcp-host", "got.games.asmodee.net", "Target TCP Host")
	flag.StringVar(&TargetTCPPort, "target-tcp-port", "2445", "Target TCP Port")
	flag.Parse()

	LogFileName = fmt.Sprintf("session_%s.jsonl", time.Now().Format("20060102_150405"))
	f, err := os.Create(LogFileName)
	if err != nil {
		log.Fatal(err)
	}
	LogFile = f
	defer LogFile.Close()

	fmt.Printf("Starting Proxy...\nLog: %s\nHTTP Target: %s\nTCP Target Default: %s:%s\n", LogFileName, TargetHTTPHost, TargetTCPHost, TargetTCPPort)

	// Start HTTP Proxy
	go startHTTPProxy()

	// Start TCP Proxy
	startTCPProxy()
}

// --- HTTP Proxy ---

func startHTTPProxy() {
	http.HandleFunc("/", handleHTTPRequest)
	fmt.Println("[HTTP] Listening on " + ProxyHTTPPort)
	log.Fatal(http.ListenAndServe(ProxyHTTPPort, nil))
}

func handleHTTPRequest(w http.ResponseWriter, r *http.Request) {
	// Log Request
	bodyBytes, _ := ioutil.ReadAll(r.Body)
	r.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes)) // Restore for forwarding

	logJSON(LogEntry{
		Protocol: "HTTP",
		Type:     "REQUEST",
		Summary:  fmt.Sprintf("%s %s", r.Method, r.URL.Path),
		Data: map[string]interface{}{
			"headers": r.Header,
			"body":    string(bodyBytes),
		},
	})

	// Forward Request
	targetURL := TargetHTTPHost + r.URL.Path
	if r.URL.RawQuery != "" {
		targetURL += "?" + r.URL.RawQuery
	}

	proxyReq, err := http.NewRequest(r.Method, targetURL, bytes.NewBuffer(bodyBytes))
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	// Copy headers
	for name, values := range r.Header {
		for _, value := range values {
			proxyReq.Header.Add(name, value)
		}
	}
	// Remove Hop-by-hop headers
	proxyReq.Header.Del("Host") // Let http client set it
	// Might need to set Host header manually if target expects it
	// proxyReq.Host = ...

	client := &http.Client{}
	resp, err := client.Do(proxyReq)
	if err != nil {
		// Log Error
		logJSON(LogEntry{
			Protocol: "HTTP",
			Type:     "ERROR",
			Summary:  "Forwarding Failed",
			Data:     err.Error(),
		})
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	respBody, _ := ioutil.ReadAll(resp.Body)

	// Log Response
	logJSON(LogEntry{
		Protocol: "HTTP",
		Type:     "RESPONSE",
		Summary:  fmt.Sprintf("%s %d", r.URL.Path, resp.StatusCode),
		Data: map[string]interface{}{
			"status":  resp.StatusCode,
			"headers": resp.Header,
			"body":    string(respBody),
		},
	})

	// Sniff for TCP Info
	// Look for "server": { "host": ... } or "HostName": ...
	sniffTCPInfo(respBody)

	// Write Response
	for name, values := range resp.Header {
		for _, value := range values {
			w.Header().Add(name, value)
		}
	}
	w.WriteHeader(resp.StatusCode)
	w.Write(respBody)
}

func sniffTCPInfo(body []byte) {
	// Search for "host":"..." and "port":... pattern
	s := string(body)
	if strings.Contains(s, "\"host\"") && strings.Contains(s, "\"port\"") {
		idxHost := strings.Index(s, "\"host\"")
		if idxHost != -1 {
			// Find value
			startQuote := strings.Index(s[idxHost:], ":") + idxHost + 1
			startQuote = strings.Index(s[startQuote:], "\"") + startQuote
			endQuote := strings.Index(s[startQuote+1:], "\"") + startQuote + 1

			if startQuote != -1 && endQuote != -1 {
				extractedHost := s[startQuote+1 : endQuote]
				if extractedHost != "" {
					fmt.Printf("[PROXY] Sniffed TCP Host: %s\n", extractedHost)
					TargetTCPHost = extractedHost
				}
			}
		}

		idxPort := strings.Index(s, "\"port\"")
		if idxPort != -1 {
			// Find value (number)
			colon := strings.Index(s[idxPort:], ":") + idxPort
			// Find first digit after colon
			startNum := -1
			endNum := -1
			for i := colon + 1; i < len(s); i++ {
				c := s[i]
				if c >= '0' && c <= '9' {
					if startNum == -1 {
						startNum = i
					}
				} else if startNum != -1 {
					// End of number
					endNum = i
					break
				}
			}
			if startNum != -1 {
				if endNum == -1 {
					endNum = len(s)
				}
				extractedPort := s[startNum:endNum]
				fmt.Printf("[PROXY] Sniffed TCP Port: %s\n", extractedPort)
				TargetTCPPort = extractedPort
			}
		}
	}
}

func findHostPortRecursive(obj interface{}) {
	switch v := obj.(type) {
	case map[string]interface{}:
		if h, ok := v["host"].(string); ok {
			TargetTCPHost = h
			fmt.Printf("[PROXY] Sniffed TCP Host: %s\n", h)
		}
		if p, ok := v["port"].(float64); ok { // JSON numbers are float64 in generic interface
			TargetTCPPort = fmt.Sprintf("%d", int(p))
			fmt.Printf("[PROXY] Sniffed TCP Port: %s\n", TargetTCPPort)
		}
		for _, val := range v {
			findHostPortRecursive(val)
		}
	case []interface{}:
		for _, val := range v {
			findHostPortRecursive(val)
		}
	}
}

// --- TCP Proxy ---

func startTCPProxy() {
	cert, err := generateSelfSignedCert()
	if err != nil {
		log.Fatal(err)
	}

	listener, err := tls.Listen("tcp", ProxyTCPPort, cert)
	if err != nil {
		log.Fatalf("TCP Listen failed: %v", err)
	}
	fmt.Println("[TCP] Listening on " + ProxyTCPPort)

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Println("TCP Accept error:", err)
			continue
		}
		go handleTCPConn(conn)
	}
}

func handleTCPConn(clientConn net.Conn) {
	defer clientConn.Close()

	// Connect to Real Server
	targetAddr := fmt.Sprintf("%s:%s", TargetTCPHost, TargetTCPPort)
	fmt.Printf("[TCP] New Connection. Forwarding to %s\n", targetAddr)

	// Use TLS to connect to Real Server
	serverConn, err := tls.Dial("tcp", targetAddr, &tls.Config{
		InsecureSkipVerify: true, // We assume official server cert is fine but we skip verify to avoid setup
	})
	if err != nil {
		log.Printf("Failed to dial target %s: %v\n", targetAddr, err)
		return
	}
	defer serverConn.Close()

	// Pipe
	var wg sync.WaitGroup
	wg.Add(2)

	// Client -> Server
	go func() {
		defer wg.Done()
		interceptAndLogTCP(clientConn, serverConn, "CLIENT->SERVER")
	}()

	// Server -> Client
	go func() {
		defer wg.Done()
		interceptAndLogTCP(serverConn, clientConn, "SERVER->CLIENT")
	}()

	wg.Wait()
}

func interceptAndLogTCP(src, dst net.Conn, direction string) {
	// We need to parse frames to log them effectively?
	// The protocol is: 4 bytes length + Payload.
	// Payload: Varint PacketID, Varint RequestID, Payload.

	// We will try to read in a buffer loop, parse length prefix, log, then forward.
	// But simply copying stream `io.Copy` is safer for latency.
	// However, we want to LOG. So we MUST read packet by packet.

	header := make([]byte, 4)
	for {
		// Read Head
		_, err := io.ReadFull(src, header)
		if err != nil {
			break
		}
		length := binary.BigEndian.Uint32(header)

		// Read Body
		body := make([]byte, length)
		_, err = io.ReadFull(src, body)
		if err != nil {
			break
		}

		// Write to Dst
		// Write header + body
		dst.Write(header)
		dst.Write(body)

		// Log
		parseAndLogPacket(direction, body)
	}
}

func parseAndLogPacket(direction string, data []byte) {
	// Basic Decode of Packet Wrapper
	// Packet ID (Field 1, Varint)
	// Payload (Field 3, Bytes) -> Message

	r := bytes.NewReader(data)

	packetID := int64(-1)
	var payloadBytes []byte

	for {
		tag, err := binary.ReadUvarint(r)
		if err != nil {
			break
		}
		fieldNum := tag >> 3
		wireType := tag & 7

		if fieldNum == 1 && wireType == 0 {
			packetID, _ = binary.ReadVarint(r) // ReadVarint for int64
		} else if fieldNum == 3 && wireType == 2 {
			l, _ := binary.ReadUvarint(r)
			payloadBytes = make([]byte, l)
			r.Read(payloadBytes)
		} else {
			// Skip
			skip(r, wireType)
		}
	}

	// If we found payload, parse Message
	requestID := int64(-1)
	if payloadBytes != nil {
		pr := bytes.NewReader(payloadBytes)
		for {
			tag, err := binary.ReadUvarint(pr)
			if err != nil {
				break
			}
			fieldNum := tag >> 3
			wireType := tag & 7

			if fieldNum == 1 && wireType == 0 { // Message.request_number
				requestID, _ = binary.ReadVarint(pr)
			} else {
				skip(pr, wireType)
			}
		}
	}

	logJSON(LogEntry{
		Protocol: "TCP",
		Type:     direction,
		Summary:  fmt.Sprintf("Packet %d | Req %d | Len %d", packetID, requestID, len(data)),
		Data: map[string]interface{}{
			"packet_id":  packetID,
			"request_id": requestID,
			"hex":        hex.EncodeToString(data),
		},
	})
}

func skip(r *bytes.Reader, wireType uint64) {
	switch wireType {
	case 0:
		binary.ReadUvarint(r)
	case 1:
		r.Seek(8, io.SeekCurrent)
	case 2:
		l, _ := binary.ReadUvarint(r)
		r.Seek(int64(l), io.SeekCurrent)
	case 5:
		r.Seek(4, io.SeekCurrent)
	}
}

// --- Cert Gen (Copied) ---
func generateSelfSignedCert() (*tls.Config, error) {
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, err
	}
	template := x509.Certificate{
		SerialNumber:          big.NewInt(1),
		Subject:               pkix.Name{Organization: []string{"Proxy"}},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour * 24),
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		IPAddresses:           []net.IP{net.ParseIP("127.0.0.1")},
		DNSNames:              []string{"localhost"},
	}
	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		return nil, err
	}
	return &tls.Config{Certificates: []tls.Certificate{{Certificate: [][]byte{derBytes}, PrivateKey: priv}}}, nil
}