package ntlmssp

import (
	"encoding/binary"
	"fmt"
	"strings"
)

// Debugging flag
var DebugMode = true

// PrintDebug logs debug messages when DebugMode is enabled
func PrintDebug(format string, args ...interface{}) {
	if DebugMode {
		fmt.Printf(format+"\n", args...)
	}
}

// DecodeNTLMMessage decodes NTLM messages and prints details
func DecodeNTLMMessage(blob []byte) {
	if len(blob) < 12 {
		PrintDebug("Invalid NTLM message (too short)")
		return
	}

	if string(blob[:8]) != "NTLMSSP\x00" {
		PrintDebug("Invalid NTLM signature")
		return
	}

	msgType := binary.LittleEndian.Uint32(blob[8:12])
	switch msgType {
	case 1:
		DecodeType1Message(blob)
	case 2:
		DecodeType2Message(blob)
	case 3:
		DecodeType3Message(blob)
	default:
		PrintDebug("Unknown NTLM message type: %d", msgType)
	}
}

// DecodeType1Message prints details of an NTLM Type 1 message
func DecodeType1Message(blob []byte) {
	PrintDebug("==== Type1 ----")
	PrintDebug("Signature: NTLMSSP")
	PrintDebug("Type: 1")

	if len(blob) < 32 {
		PrintDebug("Invalid NTLM Type 1 message")
		return
	}

	flags := binary.LittleEndian.Uint32(blob[12:16])
	PrintDebug("Flags: %08X", flags)
	PrintDebug(DecodeFlags(flags))

	domainLen := binary.LittleEndian.Uint16(blob[16:18])
	domainMaxLen := binary.LittleEndian.Uint16(blob[18:20])
	domainOffset := binary.LittleEndian.Uint32(blob[20:24])

	workstationLen := binary.LittleEndian.Uint16(blob[24:26])
	workstationMaxLen := binary.LittleEndian.Uint16(blob[26:28])
	workstationOffset := binary.LittleEndian.Uint32(blob[28:32])

	if domainMaxLen > 0 && int(domainOffset+uint32(domainLen)) <= len(blob) {
		domain := string(blob[domainOffset : domainOffset+uint32(domainLen)])
		PrintDebug("Domain: %s", domain)
	}

	if workstationMaxLen > 0 && int(workstationOffset+uint32(workstationLen)) <= len(blob) {
		workstation := string(blob[workstationOffset : workstationOffset+uint32(workstationLen)])
		PrintDebug("Workstation: %s", workstation)
	}
}

// DecodeType2Message prints details of an NTLM Type 2 message
func DecodeType2Message(blob []byte) {
	PrintDebug("==== Type2 ----")
	PrintDebug("Signature: NTLMSSP")
	PrintDebug("Type: 2")

	if len(blob) < 48 {
		PrintDebug("Invalid NTLM Type 2 message")
		return
	}

	flags := binary.LittleEndian.Uint32(blob[20:24])
	PrintDebug("Flags: %08X", flags)
	PrintDebug(DecodeFlags(flags))

	challenge := blob[24:32]
	PrintDebug("Challenge: %X", challenge)

	context := blob[40:48]
	PrintDebug("Context: %X:%X", context[:4], context[4:])
}

// DecodeType3Message prints details of an NTLM Type 3 message
func DecodeType3Message(blob []byte) {
	PrintDebug("==== Type3 ----")
	PrintDebug("Signature: NTLMSSP")
	PrintDebug("Type: 3")

	if len(blob) < 64 {
		PrintDebug("Invalid NTLM Type 3 message")
		return
	}

	flags := binary.LittleEndian.Uint32(blob[60:64])
	PrintDebug("Flags: %08X", flags)
	PrintDebug(DecodeFlags(flags))
}

// DecodeFlags returns a formatted string listing NTLM flag names
func DecodeFlags(flags uint32) string {
	var flagStrings []string

	flagMap := map[uint32]string{
		0x00000001: "NTLMSSP_NEGOTIATE_UNICODE",
		0x00000002: "NTLMSSP_NEGOTIATE_OEM",
		0x00000004: "NTLMSSP_REQUEST_TARGET",
		0x00000010: "NTLMSSP_NEGOTIATE_SIGN",
		0x00000020: "NTLMSSP_NEGOTIATE_SEAL",
		0x00000040: "NTLMSSP_NEGOTIATE_DATAGRAM",
		0x00000080: "NTLMSSP_NEGOTIATE_LM_KEY",
		0x00000100: "NTLMSSP_NEGOTIATE_NETWARE",
		0x00000200: "NTLMSSP_NEGOTIATE_NTLM",
		0x00000800: "NTLMSSP_NEGOTIATE_OEM_DOMAIN_SUPPLIED",
		0x00001000: "NTLMSSP_NEGOTIATE_OEM_WORKSTATION_SUPPLIED",
		0x00002000: "NTLMSSP_NEGOTIATE_ALWAYS_SIGN",
		0x00020000: "NTLMSSP_NEGOTIATE_TARGET_INFO",
		0x00040000: "NTLMSSP_REQUEST_NON_NT_SESSION_KEY",
		0x00080000: "NTLMSSP_NEGOTIATE_EXTENDED_SESSIONSECURITY",
		0x00100000: "NTLMSSP_NEGOTIATE_IDENTIFY",
		0x00200000: "NTLMSSP_REQUEST_TARGET",
		0x00800000: "NTLMSSP_TARGET_TYPE_DOMAIN",
		0x01000000: "NTLMSSP_TARGET_TYPE_SERVER",
		0x02000000: "NTLMSSP_TARGET_TYPE_SHARE",
		0x08000000: "NTLMSSP_NEGOTIATE_KEY_EXCH",
		0x20000000: "NTLMSSP_NEGOTIATE_128",
		0x80000000: "NTLMSSP_NEGOTIATE_56",
	}

	for bit, name := range flagMap {
		if flags&bit != 0 {
			flagStrings = append(flagStrings, name)
		}
	}

	return strings.Join(flagStrings, "\n")
}