package smb

import (
	"bytes"
	"encoding/binary"
	"git.loafle.net/overflow/overflow_discovery/match/packet"
	"git.loafle.net/overflow/overflow_discovery/model/scaninfo"
	"strings"
)

const (
	SMB_COM_NEGOTIATE uint8 = 0x72
	SMB_SUCCESS       uint8 = 0x00
)

type netBIOS struct {
	MsgType   byte
	MsgLength [3]uint8
}

type smb struct {
	NetBios    netBIOS
	Component  [4]uint8
	SmbCommand uint8
	NtStatus   [4]uint8
	Flags      uint8
	Flags2     [2]uint8
	ProcessId  uint16
	Signature  uint64
	Reserved   uint16
	Tid        uint16
	Pid        uint16
	Uid        uint16
	Mid        uint16
	Wct        uint8
	Bcc        uint16
	Bf1        uint8
	Name1      [23]uint8
	Bf2        uint8
	Name2      [10]uint8
	Bf3        uint8
	Name3      [28]uint8
	Bf4        uint8
	Name4      [10]uint8
	Bf5        uint8
	Name5      [10]uint8
	Bf6        uint8
	Name6      [11]uint8
}

type SMBMatcher struct {
	packets []*packet.Packet
}

func NewSMBMatcher() *SMBMatcher {

	nbssMatcher := &SMBMatcher{}

	query := smb{}
	query.NetBios.MsgType = 0x00
	query.NetBios.MsgLength[2] = 0x85

	query.Component[0] = 0xff
	query.Component[1] = 'S'
	query.Component[2] = 'M'
	query.Component[3] = 'B'

	query.SmbCommand = SMB_COM_NEGOTIATE
	query.NtStatus[3] = SMB_SUCCESS
	query.Flags = 0x18
	query.Flags2[0] = 0x53
	query.Flags2[1] = 0xC8

	query.ProcessId = 0x00
	query.Signature = 0x00
	query.Reserved = 0
	query.Tid = 0
	query.Pid = 0xfeff
	query.Uid = 0
	query.Mid = 0
	query.Wct = 0
	query.Bcc = 0x0062

	query.Bf1 = 0x02
	copy(query.Name1[:], "PC NETWORK PROGRAM 1.0")

	query.Bf2 = 0x02
	copy(query.Name2[:], "LANMAN1.0")

	query.Bf3 = 0x02
	copy(query.Name3[:], "Windows for Workgroups 3.1a")

	query.Bf4 = 0x02
	copy(query.Name4[:], "LM1.2X002")

	query.Bf5 = 0x02
	copy(query.Name5[:], "LANMAN2.1")

	query.Bf6 = 0x02
	copy(query.Name6[:], "NT LM 0.12")

	writer := new(bytes.Buffer)
	binary.Write(writer, binary.LittleEndian, query)

	nbssMatcher.packets = append(nbssMatcher.packets, packet.NewPacket(writer.Bytes(), writer.Len()))

	return nbssMatcher
}

func (t *SMBMatcher) ServiceName() string {
	return "SMB"
}

func (t *SMBMatcher) PacketCount() int {
	return len(t.packets)
}
func (t *SMBMatcher) Packet(index int) *packet.Packet {
	return t.packets[index]
}

func (t *SMBMatcher) IsError(index int, packet *packet.Packet, info scaninfo.ServiceScanInfo) bool {
	return false
}

func (t *SMBMatcher) HasResponse(index int) bool {
	return false
}

func (t *SMBMatcher) IsPrePacket() bool {
	return false
}

func (t *SMBMatcher) Match(index int, packet *packet.Packet, info scaninfo.ServiceScanInfo) bool {

	if packet == nil {
		return false
	}

	reader := new(bytes.Buffer)
	reader.Write(packet.Buffer)

	s := smb{}
	if err := binary.Read(reader, binary.BigEndian, &s); err != nil {
		return false
	}

	var des [4]byte
	copy(des[1:], s.NetBios.MsgLength[:])
	packetLen := binary.BigEndian.Uint32(des[:])

	if packetLen != uint32(packet.Len-4) {
		return false
	}

	if !strings.Contains(string(s.Component[:]), "SMB") {
		return false
	}

	return true

}