package tlssniffercore

import (
	"bytes"
	"encoding/binary"
	"io"
	"net"
)

type ClientHelloMeta struct {
	ServerName        string
	LocalAddr         net.Addr
	RemoteAddr        net.Addr
	SupportedProtos   []string
	SupportedVersions []uint16
	CipherSuites      []uint16
}

type SniffResult struct {
	IsTLS       bool
	ClientHello *ClientHelloMeta
	Buffer      *bytes.Buffer
}

type Sniffer struct{}

func (s Sniffer) Sniff(conn net.Conn, maxBytes int) (SniffResult, error) {
	if maxBytes <= 0 {
		maxBytes = 64 * 1024
	}

	var buf bytes.Buffer
	limited := &io.LimitedReader{R: conn, N: int64(maxBytes)}
	meta, isTLS := sniffClientHello(limited, &buf, conn)

	out := SniffResult{
		IsTLS:  isTLS,
		Buffer: bytes.NewBuffer(append([]byte(nil), buf.Bytes()...)),
	}
	if isTLS {
		out.ClientHello = meta
	}
	return out, nil
}

func sniffClientHello(reader io.Reader, buf *bytes.Buffer, conn net.Conn) (*ClientHelloMeta, bool) {
	meta := &ClientHelloMeta{
		LocalAddr:  conn.LocalAddr(),
		RemoteAddr: conn.RemoteAddr(),
	}

	header, complete := readTLSRecordHeader(reader, buf)
	if len(header) < 3 {
		return nil, false
	}
	isTLS := header[0] == 0x16 && header[1] == 0x03
	if !isTLS {
		return nil, false
	}
	if len(header) < 5 || !complete {
		return meta, true
	}

	recordLen := int(binary.BigEndian.Uint16(header[3:5]))
	recordBody, bodyOK := readBufferedBytes(reader, buf, recordLen)
	if !bodyOK {
		return meta, true
	}
	if len(recordBody) < 4 || recordBody[0] != 0x01 {
		return nil, false
	}

	helloLen := int(recordBody[1])<<16 | int(recordBody[2])<<8 | int(recordBody[3])
	helloBytes := append([]byte(nil), recordBody[4:]...)
	for len(helloBytes) < helloLen {
		nextHeader, ok := readTLSRecordHeader(reader, buf)
		if len(nextHeader) < 5 || !ok {
			return meta, true
		}
		if nextHeader[0] != 0x16 || nextHeader[1] != 0x03 {
			return meta, true
		}
		nextLen := int(binary.BigEndian.Uint16(nextHeader[3:5]))
		nextBody, bodyOK := readBufferedBytes(reader, buf, nextLen)
		if !bodyOK {
			return meta, true
		}
		helloBytes = append(helloBytes, nextBody...)
	}

	parseClientHelloBody(meta, helloBytes[:helloLen])
	return meta, true
}

func readTLSRecordHeader(reader io.Reader, buf *bytes.Buffer) ([]byte, bool) {
	return readBufferedBytes(reader, buf, 5)
}

func readBufferedBytes(reader io.Reader, buf *bytes.Buffer, count int) ([]byte, bool) {
	if count <= 0 {
		return nil, true
	}
	tmp := make([]byte, count)
	readN, err := io.ReadFull(reader, tmp)
	if readN > 0 {
		buf.Write(tmp[:readN])
	}
	return append([]byte(nil), tmp[:readN]...), err == nil
}

func parseClientHelloBody(meta *ClientHelloMeta, body []byte) {
	if meta == nil || len(body) < 34 {
		return
	}

	offset := 2 + 32
	sessionIDLen := int(body[offset])
	offset++
	if offset+sessionIDLen > len(body) {
		return
	}
	offset += sessionIDLen

	if offset+2 > len(body) {
		return
	}
	cipherSuitesLen := int(binary.BigEndian.Uint16(body[offset : offset+2]))
	offset += 2
	if offset+cipherSuitesLen > len(body) {
		return
	}
	for index := 0; index+1 < cipherSuitesLen; index += 2 {
		meta.CipherSuites = append(meta.CipherSuites, binary.BigEndian.Uint16(body[offset+index:offset+index+2]))
	}
	offset += cipherSuitesLen

	if offset >= len(body) {
		return
	}
	compressionMethodsLen := int(body[offset])
	offset++
	if offset+compressionMethodsLen > len(body) {
		return
	}
	offset += compressionMethodsLen

	if offset+2 > len(body) {
		return
	}
	extensionsLen := int(binary.BigEndian.Uint16(body[offset : offset+2]))
	offset += 2
	if offset+extensionsLen > len(body) {
		return
	}

	parseClientHelloExtensions(meta, body[offset:offset+extensionsLen])
}

func parseClientHelloExtensions(meta *ClientHelloMeta, exts []byte) {
	for offset := 0; offset+4 <= len(exts); {
		extType := binary.BigEndian.Uint16(exts[offset : offset+2])
		extLen := int(binary.BigEndian.Uint16(exts[offset+2 : offset+4]))
		offset += 4
		if offset+extLen > len(exts) {
			return
		}
		extData := exts[offset : offset+extLen]
		offset += extLen

		switch extType {
		case 0:
			parseServerNameExtension(meta, extData)
		case 16:
			parseALPNExtension(meta, extData)
		case 43:
			parseSupportedVersionsExtension(meta, extData)
		}
	}
}

func parseServerNameExtension(meta *ClientHelloMeta, data []byte) {
	if len(data) < 2 {
		return
	}
	listLen := int(binary.BigEndian.Uint16(data[:2]))
	if listLen == 0 || 2+listLen > len(data) {
		return
	}
	list := data[2 : 2+listLen]
	for offset := 0; offset+3 <= len(list); {
		nameType := list[offset]
		nameLen := int(binary.BigEndian.Uint16(list[offset+1 : offset+3]))
		offset += 3
		if offset+nameLen > len(list) {
			return
		}
		if nameType == 0 {
			meta.ServerName = string(list[offset : offset+nameLen])
			return
		}
		offset += nameLen
	}
}

func parseALPNExtension(meta *ClientHelloMeta, data []byte) {
	if len(data) < 2 {
		return
	}
	listLen := int(binary.BigEndian.Uint16(data[:2]))
	if listLen == 0 || 2+listLen > len(data) {
		return
	}
	list := data[2 : 2+listLen]
	for offset := 0; offset < len(list); {
		nameLen := int(list[offset])
		offset++
		if offset+nameLen > len(list) {
			return
		}
		meta.SupportedProtos = append(meta.SupportedProtos, string(list[offset:offset+nameLen]))
		offset += nameLen
	}
}

func parseSupportedVersionsExtension(meta *ClientHelloMeta, data []byte) {
	if len(data) < 1 {
		return
	}
	listLen := int(data[0])
	if listLen == 0 || 1+listLen > len(data) {
		return
	}
	list := data[1 : 1+listLen]
	for offset := 0; offset+1 < len(list); offset += 2 {
		meta.SupportedVersions = append(meta.SupportedVersions, binary.BigEndian.Uint16(list[offset:offset+2]))
	}
}