package credssp

import (
	"bytes"
	"crypto/hmac"
	"crypto/md5"
	"crypto/rand"
	"crypto/rc4"
	"encoding/binary"
	"fmt"
	"strings"
	"time"

	"golang.org/x/crypto/md4"
)

const (
	ntlmNegotiateUnicode                 uint32 = 0x00000001
	ntlmRequestTarget                    uint32 = 0x00000004
	ntlmNegotiateSign                    uint32 = 0x00000010
	ntlmNegotiateSeal                    uint32 = 0x00000020
	ntlmNegotiateNTLM                    uint32 = 0x00000200
	ntlmNegotiateAlwaysSign              uint32 = 0x00008000
	ntlmNegotiateExtendedSessionSecurity uint32 = 0x00080000
	ntlmNegotiateTargetInfo              uint32 = 0x00800000
	ntlmNegotiateVersion                 uint32 = 0x02000000
	ntlmNegotiate128                     uint32 = 0x20000000
	ntlmNegotiateKeyExch                 uint32 = 0x40000000
	ntlmNegotiate56                      uint32 = 0x80000000

	msvAvEOL            uint16 = 0x0000
	msvAvNbComputerName uint16 = 0x0001
	msvAvNbDomainName   uint16 = 0x0002
	msvAvFlags          uint16 = 0x0006
	msvAvTimestamp      uint16 = 0x0007
	msvChannelBindings  uint16 = 0x000A
)

const ntlmClientFlags = ntlmNegotiateUnicode |
	ntlmRequestTarget |
	ntlmNegotiateSign |
	ntlmNegotiateSeal |
	ntlmNegotiateNTLM |
	ntlmNegotiateAlwaysSign |
	ntlmNegotiateExtendedSessionSecurity |
	ntlmNegotiateTargetInfo |
	ntlmNegotiateVersion |
	ntlmNegotiate128 |
	ntlmNegotiateKeyExch

type ntlmChallenge struct {
	Raw             []byte
	Flags           uint32
	TargetName      string
	ServerChallenge []byte
	TargetInfo      []byte
	Timestamp       []byte
}

type ntlmOptions struct {
	ForceMIC        bool
	DisableKeyExch  bool
	ChannelBindings []byte
	OmitVersion     bool
}

type ntlmSecurity struct {
	clientSigningKey []byte
	serverSigningKey []byte
	clientSealing    *rc4.Cipher
	serverSealing    *rc4.Cipher
	seq              uint32
}

func buildNegotiate() []byte {
	buf := bytes.NewBuffer(nil)
	buf.Write([]byte("NTLMSSP\x00"))
	writeU32(buf, 1)
	writeU32(buf, ntlmClientFlags)
	writeSecBuffer(buf, 0, 0)
	writeSecBuffer(buf, 0, 0)
	buf.Write(ntlmVersion())
	return buf.Bytes()
}

func parseChallenge(data []byte) (*ntlmChallenge, error) {
	if len(data) < 48 || string(data[:8]) != "NTLMSSP\x00" || binary.LittleEndian.Uint32(data[8:12]) != 2 {
		return nil, fmt.Errorf("invalid NTLM challenge")
	}
	targetNameLen := int(binary.LittleEndian.Uint16(data[12:14]))
	targetNameOffset := int(binary.LittleEndian.Uint32(data[16:20]))
	targetInfoLen := int(binary.LittleEndian.Uint16(data[40:42]))
	targetInfoOffset := int(binary.LittleEndian.Uint32(data[44:48]))
	if targetInfoOffset < 0 || targetInfoOffset+targetInfoLen > len(data) {
		return nil, fmt.Errorf("invalid NTLM target info")
	}
	targetName := ""
	if targetNameOffset >= 0 && targetNameOffset+targetNameLen <= len(data) {
		targetName = utf16leString(data[targetNameOffset : targetNameOffset+targetNameLen])
	}
	targetInfo := data[targetInfoOffset : targetInfoOffset+targetInfoLen]
	return &ntlmChallenge{
		Raw:             append([]byte(nil), data...),
		Flags:           binary.LittleEndian.Uint32(data[20:24]),
		TargetName:      targetName,
		ServerChallenge: append([]byte(nil), data[24:32]...),
		TargetInfo:      append([]byte(nil), targetInfo...),
		Timestamp:       findAVPair(targetInfo, msvAvTimestamp),
	}, nil
}

func buildAuthenticate(ch *ntlmChallenge, domain, user, password, workstation string, negotiate []byte, opts ntlmOptions) ([]byte, *ntlmSecurity, error) {
	if domain == "" {
		domain = ch.TargetName
	}
	if domain == "" {
		domain = utf16leString(findAVPair(ch.TargetInfo, msvAvNbDomainName))
	}
	respKey := ntowfv2(password, user, domain)
	clientChallenge := randomBytes(8)
	timestamp := ch.Timestamp
	if len(timestamp) == 0 {
		timestamp = windowsFileTime(time.Now())
	}

	blob := bytes.NewBuffer(nil)
	blob.Write([]byte{0x01, 0x01, 0, 0, 0, 0, 0, 0})
	blob.Write(timestamp)
	blob.Write(clientChallenge)
	blob.Write([]byte{0, 0, 0, 0})
	targetInfo := ch.TargetInfo
	if len(opts.ChannelBindings) == 16 {
		targetInfo = withChannelBindings(ch.TargetInfo, opts.ChannelBindings)
	}
	blob.Write(targetInfo)
	blob.Write([]byte{0, 0, 0, 0})

	ntProof := hmacMD5(respKey, append(append([]byte(nil), ch.ServerChallenge...), blob.Bytes()...))
	ntResp := append(ntProof, blob.Bytes()...)
	lmResp := append(hmacMD5(respKey, append(append([]byte(nil), ch.ServerChallenge...), clientChallenge...)), clientChallenge...)
	sessionBaseKey := hmacMD5(respKey, ntProof)
	exportedSessionKey := randomBytes(16)
	encryptedSessionKey := rc4k(sessionBaseKey, exportedSessionKey)

	domainBytes := utf16le(domain)
	userBytes := utf16le(user)
	workstationBytes := utf16le(workstation)
	flags := ch.Flags & ntlmClientFlags
	flags |= ntlmNegotiateUnicode | ntlmNegotiateExtendedSessionSecurity | ntlmNegotiateNTLM
	if opts.DisableKeyExch {
		flags &^= ntlmNegotiateKeyExch
		exportedSessionKey = sessionBaseKey
		encryptedSessionKey = nil
	}

	includeMIC := opts.ForceMIC || requiresMIC(ch.TargetInfo)
	baseLen := 72
	if opts.OmitVersion {
		baseLen = 64
	}
	if includeMIC {
		baseLen += 16
	}
	payload := bytes.NewBuffer(nil)
	offset := uint32(baseLen)
	lmOffset := offset
	payload.Write(lmResp)
	offset += uint32(len(lmResp))
	ntOffset := offset
	payload.Write(ntResp)
	offset += uint32(len(ntResp))
	domainOffset := offset
	payload.Write(domainBytes)
	offset += uint32(len(domainBytes))
	userOffset := offset
	payload.Write(userBytes)
	offset += uint32(len(userBytes))
	workstationOffset := offset
	payload.Write(workstationBytes)
	offset += uint32(len(workstationBytes))
	sessionOffset := offset
	if len(encryptedSessionKey) > 0 {
		payload.Write(encryptedSessionKey)
	}

	msg := bytes.NewBuffer(nil)
	msg.Write([]byte("NTLMSSP\x00"))
	writeU32(msg, 3)
	writeSecBuffer(msg, uint16(len(lmResp)), lmOffset)
	writeSecBuffer(msg, uint16(len(ntResp)), ntOffset)
	writeSecBuffer(msg, uint16(len(domainBytes)), domainOffset)
	writeSecBuffer(msg, uint16(len(userBytes)), userOffset)
	writeSecBuffer(msg, uint16(len(workstationBytes)), workstationOffset)
	writeSecBuffer(msg, uint16(len(encryptedSessionKey)), sessionOffset)
	writeU32(msg, flags)
	if !opts.OmitVersion {
		msg.Write(ntlmVersion())
	}
	micPos := -1
	if includeMIC {
		micPos = msg.Len()
		msg.Write(make([]byte, 16))
	}
	msg.Write(payload.Bytes())

	if includeMIC {
		micInput := bytes.NewBuffer(nil)
		micInput.Write(negotiate)
		micInput.Write(ch.Raw)
		micInput.Write(msg.Bytes())
		mic := hmacMD5(exportedSessionKey, micInput.Bytes())
		copy(msg.Bytes()[micPos:micPos+16], mic)
	}

	clientSeal, _ := rc4.NewCipher(md5sum(append(exportedSessionKey, []byte("session key to client-to-server sealing key magic constant\x00")...)))
	serverSeal, _ := rc4.NewCipher(md5sum(append(exportedSessionKey, []byte("session key to server-to-client sealing key magic constant\x00")...)))
	sec := &ntlmSecurity{
		clientSigningKey: md5sum(append(exportedSessionKey, []byte("session key to client-to-server signing key magic constant\x00")...)),
		serverSigningKey: md5sum(append(exportedSessionKey, []byte("session key to server-to-client signing key magic constant\x00")...)),
		clientSealing:    clientSeal,
		serverSealing:    serverSeal,
	}
	return msg.Bytes(), sec, nil
}

func (s *ntlmSecurity) encrypt(data []byte) []byte {
	sealed := make([]byte, len(data))
	s.clientSealing.XORKeyStream(sealed, data)
	macInput := bytes.NewBuffer(nil)
	writeU32(macInput, s.seq)
	macInput.Write(data)
	checksum := hmacMD5(s.clientSigningKey, macInput.Bytes())[:8]
	encryptedChecksum := make([]byte, 8)
	s.clientSealing.XORKeyStream(encryptedChecksum, checksum)
	out := bytes.NewBuffer(nil)
	writeU32(out, 1)
	out.Write(encryptedChecksum)
	writeU32(out, s.seq)
	out.Write(sealed)
	s.seq++
	return out.Bytes()
}

func ntowfv2(password, user, domain string) []byte {
	h := md4.New()
	h.Write(utf16le(password))
	return hmacMD5(h.Sum(nil), utf16le(strings.ToUpper(user)+domain))
}

func hmacMD5(key, data []byte) []byte {
	h := hmac.New(md5.New, key)
	h.Write(data)
	return h.Sum(nil)
}

func md5sum(data []byte) []byte {
	sum := md5.Sum(data)
	return sum[:]
}

func rc4k(key, data []byte) []byte {
	c, _ := rc4.NewCipher(key)
	out := make([]byte, len(data))
	c.XORKeyStream(out, data)
	return out
}

func randomBytes(n int) []byte {
	buf := make([]byte, n)
	_, _ = rand.Read(buf)
	return buf
}

func windowsFileTime(t time.Time) []byte {
	v := uint64(t.UnixNano()/100) + 116444736000000000
	out := make([]byte, 8)
	binary.LittleEndian.PutUint64(out, v)
	return out
}

func findAVPair(data []byte, id uint16) []byte {
	for len(data) >= 4 {
		avID := binary.LittleEndian.Uint16(data[0:2])
		avLen := int(binary.LittleEndian.Uint16(data[2:4]))
		data = data[4:]
		if avLen > len(data) {
			return nil
		}
		if avID == msvAvEOL {
			return nil
		}
		if avID == id {
			return append([]byte(nil), data[:avLen]...)
		}
		data = data[avLen:]
	}
	return nil
}

func requiresMIC(targetInfo []byte) bool {
	flags := findAVPair(targetInfo, msvAvFlags)
	if len(flags) < 4 {
		return false
	}
	return binary.LittleEndian.Uint32(flags)&0x00000002 != 0
}

func withChannelBindings(targetInfo, hash []byte) []byte {
	out := bytes.NewBuffer(nil)
	inserted := false
	for len(targetInfo) >= 4 {
		avID := binary.LittleEndian.Uint16(targetInfo[0:2])
		avLen := int(binary.LittleEndian.Uint16(targetInfo[2:4]))
		if 4+avLen > len(targetInfo) {
			return targetInfo
		}
		if avID == msvAvEOL {
			if !inserted {
				writeU16(out, msvChannelBindings)
				writeU16(out, 16)
				out.Write(hash[:16])
			}
			out.Write(targetInfo[:4])
			return out.Bytes()
		}
		if avID == msvChannelBindings {
			writeU16(out, msvChannelBindings)
			writeU16(out, 16)
			out.Write(hash[:16])
			inserted = true
		} else {
			out.Write(targetInfo[:4+avLen])
		}
		targetInfo = targetInfo[4+avLen:]
	}
	if !inserted {
		writeU16(out, msvChannelBindings)
		writeU16(out, 16)
		out.Write(hash[:16])
	}
	out.Write([]byte{0, 0, 0, 0})
	return out.Bytes()
}

func writeSecBuffer(buf *bytes.Buffer, length uint16, offset uint32) {
	writeU16(buf, length)
	writeU16(buf, length)
	writeU32(buf, offset)
}

func writeU16(buf *bytes.Buffer, v uint16) {
	_ = binary.Write(buf, binary.LittleEndian, v)
}

func writeU32(buf *bytes.Buffer, v uint32) {
	_ = binary.Write(buf, binary.LittleEndian, v)
}

func ntlmVersion() []byte {
	out := []byte{10, 0, 0, 0, 0, 0, 0, 15}
	binary.LittleEndian.PutUint16(out[2:4], 19045)
	return out
}

func utf16le(s string) []byte {
	out := bytes.NewBuffer(nil)
	for _, r := range s {
		if r > 0xffff {
			r = '?'
		}
		writeU16(out, uint16(r))
	}
	return out.Bytes()
}

func utf16leString(data []byte) string {
	runes := make([]rune, 0, len(data)/2)
	for i := 0; i+1 < len(data); i += 2 {
		v := binary.LittleEndian.Uint16(data[i : i+2])
		if v == 0 {
			continue
		}
		runes = append(runes, rune(v))
	}
	return string(runes)
}