// Command usbmux-proxy ist ein Man-in-the-Middle Unix-socket proxy zwischen
// einem Client (z.B. TechLockdown) und Apple's `/var/run/usbmuxd`-Daemon.
//
// Architektur:
//
//	Client --[unix socket /tmp/usbmux-proxy]--> proxy --[unix socket /var/run/usbmuxd]--> usbmuxd
//	                                              |
//	                                              v
//	                                   /tmp/usbmux-capture-*.log
//
// Usage:
//
//	# Terminal 1 — start proxy
//	./bin/rebreak-usbmux-proxy
//
//	# Terminal 2 — run TL through proxy (kein sudo nötig!)
//	USBMUXD_SOCKET_ADDRESS=unix:///tmp/usbmux-proxy \
//	  open ~/Downloads/TechLockdown-supervise-mac-arm64.app
//
// Log-Output: hex+ascii dump für jeden gesendeten/empfangenen byte.
// Plus: usbmux-frame-parsing (16-byte header + payload), plist-detection.
package main

import (
	"bytes"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"howett.net/plist"
)

const (
	defaultProxyPath = "/tmp/usbmux-proxy"
	defaultRealPath  = "/var/run/usbmuxd"
)

var (
	flagProxyPath = flag.String("proxy", defaultProxyPath, "unix-socket-pfad für proxy")
	flagRealPath  = flag.String("real", defaultRealPath, "unix-socket-pfad zum echten usbmuxd")
	flagLogPath   = flag.String("log", "", "log-output (default: /tmp/usbmux-capture-TIMESTAMP.log)")
	flagQuiet     = flag.Bool("q", false, "kein stdout output, nur file-log")

	sessionCounter atomic.Uint64
)

func main() {
	flag.Parse()

	logPath := *flagLogPath
	if logPath == "" {
		logPath = fmt.Sprintf("/tmp/usbmux-capture-%s.log", time.Now().Format("20060102-150405"))
	}
	logFile, err := os.Create(logPath)
	if err != nil {
		fmt.Fprintln(os.Stderr, "log create:", err)
		os.Exit(1)
	}
	defer logFile.Close()

	logger := &mitm{
		out:   logFile,
		quiet: *flagQuiet,
	}
	logger.printf("== usbmux-proxy starting — proxy=%s → real=%s, log=%s\n", *flagProxyPath, *flagRealPath, logPath)

	// Cleanup any stale socket
	os.Remove(*flagProxyPath)

	listener, err := net.Listen("unix", *flagProxyPath)
	if err != nil {
		fmt.Fprintln(os.Stderr, "listen:", err)
		os.Exit(1)
	}
	defer listener.Close()
	defer os.Remove(*flagProxyPath)

	logger.printf("== listening — point USBMUXD_SOCKET_ADDRESS=unix://%s\n", *flagProxyPath)

	for {
		clientConn, err := listener.Accept()
		if err != nil {
			logger.printf("accept error: %v\n", err)
			continue
		}
		sessionID := sessionCounter.Add(1)
		go handleSession(sessionID, clientConn, logger)
	}
}

type mitm struct {
	mu    sync.Mutex
	out   *os.File
	quiet bool
}

func (m *mitm) printf(format string, args ...interface{}) {
	m.mu.Lock()
	defer m.mu.Unlock()
	s := fmt.Sprintf(format, args...)
	m.out.WriteString(s)
	if !m.quiet {
		os.Stderr.WriteString(s)
	}
}

func handleSession(sessionID uint64, client net.Conn, logger *mitm) {
	defer client.Close()
	logger.printf("\n=== SESSION %d START (client connected) ===\n", sessionID)

	daemon, err := net.Dial("unix", *flagRealPath)
	if err != nil {
		logger.printf("[session %d] dial daemon: %v\n", sessionID, err)
		return
	}
	defer daemon.Close()

	var wg sync.WaitGroup
	wg.Add(2)

	// Client → Daemon
	go func() {
		defer wg.Done()
		pipeWithLog(sessionID, "C→D", client, daemon, logger)
		daemon.Close() // signal other side
	}()

	// Daemon → Client
	go func() {
		defer wg.Done()
		pipeWithLog(sessionID, "D→C", daemon, client, logger)
		client.Close()
	}()

	wg.Wait()
	logger.printf("=== SESSION %d END ===\n", sessionID)
}

// pipeWithLog liest from src + writes to dst, logging each chunk.
func pipeWithLog(sessionID uint64, dir string, src, dst net.Conn, logger *mitm) {
	buf := make([]byte, 64*1024)
	for {
		n, err := src.Read(buf)
		if n > 0 {
			chunk := buf[:n]
			logger.printf("\n--- [session %d] %s | %d bytes | %s ---\n", sessionID, dir, n, time.Now().Format("15:04:05.000"))
			dumpChunk(chunk, logger)
			if _, werr := dst.Write(chunk); werr != nil {
				logger.printf("write err: %v\n", werr)
				return
			}
		}
		if err != nil {
			if err != io.EOF {
				logger.printf("read err: %v\n", err)
			}
			return
		}
	}
}

// dumpChunk format-aware: tries to detect usbmux-frame, plist, etc.
// Falls back to hex+ascii.
func dumpChunk(data []byte, logger *mitm) {
	// 1) Try usbmux-frame parsing (16-byte header + payload)
	if len(data) >= 16 {
		length := binary.LittleEndian.Uint32(data[0:4])
		version := binary.LittleEndian.Uint32(data[4:8])
		request := binary.LittleEndian.Uint32(data[8:12])
		tag := binary.LittleEndian.Uint32(data[12:16])
		if length > 16 && int(length) <= len(data) && version <= 2 && request <= 100 {
			payload := data[16:length]
			logger.printf("  [usbmux-frame] len=%d ver=%d req=%d tag=%d payload=%d bytes\n",
				length, version, request, tag, len(payload))
			tryParsePlist(payload, logger)
			// remaining bytes
			if int(length) < len(data) {
				logger.printf("  [+%d bytes after frame]\n", len(data)-int(length))
				dumpHex(data[length:], logger, 4)
			}
			return
		}
	}

	// 2) Try mobilebackup2 / lockdown 4-byte length-prefix + plist
	if len(data) >= 4 {
		length := binary.BigEndian.Uint32(data[0:4])
		if int(length)+4 <= len(data) && length > 0 && length < 1024*1024 {
			payload := data[4 : 4+length]
			logger.printf("  [length-prefixed] len=%d payload=%d bytes\n", length, len(payload))
			tryParsePlist(payload, logger)
			if int(length)+4 < len(data) {
				logger.printf("  [+%d bytes after frame]\n", len(data)-int(length)-4)
				dumpHex(data[int(length)+4:], logger, 4)
			}
			return
		}
	}

	// 3) Raw hex dump
	dumpHex(data, logger, 4)
}

// tryParsePlist — try XML or binary plist parse, log readable form.
func tryParsePlist(data []byte, logger *mitm) {
	if bytes.HasPrefix(data, []byte("<?xml")) || bytes.HasPrefix(data, []byte("bplist00")) {
		var parsed interface{}
		if _, err := plist.Unmarshal(data, &parsed); err == nil {
			// Re-encode as XML for readability
			var buf bytes.Buffer
			enc := plist.NewEncoderForFormat(&buf, plist.XMLFormat)
			enc.Indent("  ")
			if enc.Encode(parsed) == nil {
				logger.printf("  [plist parsed]:\n%s\n", indent(buf.String(), "    "))
				return
			}
		}
	}
	dumpHex(data, logger, 4)
}

// dumpHex — print hex+ascii dump.
func dumpHex(data []byte, logger *mitm, indentSpaces int) {
	const lineLen = 16
	prefix := strings.Repeat(" ", indentSpaces)
	for i := 0; i < len(data); i += lineLen {
		end := i + lineLen
		if end > len(data) {
			end = len(data)
		}
		hex := ""
		ascii := ""
		for j := i; j < end; j++ {
			hex += fmt.Sprintf("%02x ", data[j])
			if data[j] >= 32 && data[j] < 127 {
				ascii += string(data[j])
			} else {
				ascii += "."
			}
		}
		hex += strings.Repeat("   ", lineLen-(end-i))
		logger.printf("%s%04x: %s | %s\n", prefix, i, hex, ascii)
	}
}

func indent(s, pad string) string {
	lines := strings.Split(s, "\n")
	for i, line := range lines {
		if line != "" {
			lines[i] = pad + line
		}
	}
	return strings.Join(lines, "\n")
}