package wmi

import (
	"bytes"
	"encoding/binary"

	csm "git.loafle.net/commons/service_matcher-go"
)

const (
	PDU_BIND     = 11
	PDU_BIND_ACK = 12
	PDU_REQ      = 0
	PDU_RESP     = 2

	WMI_CALL_ID_1 = 0x95
	WMI_CALL_ID_2 = 0x96
)

type WMIMatcher struct {
	csm.Matchers
}

func (w *WMIMatcher) Key() string {
	return "WMI"
}

func (w *WMIMatcher) String() string {
	return "WMI"
}

func (w *WMIMatcher) IsPrePacket() bool {
	return false
}

func (w *WMIMatcher) HasResponse(index int) bool {
	return true
}

func (w *WMIMatcher) IsError(info csm.MatchInfo, index int, packet *csm.Packet) bool {
	return false
}

func (w *WMIMatcher) Match(info csm.MatchInfo, index int, packet *csm.Packet) bool {

	if packet == nil {
		return false
	}

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

	wmiRecv := DCERPC_DEFAULT{}

	binary.Read(buf, binary.LittleEndian, &wmiRecv)

	switch index {
	case 0:
		if wmiRecv.Call_id != WMI_CALL_ID_1 {
			return false
		}

		if wmiRecv.Ptype != PDU_BIND_ACK {
			return false
		}

		return true
	case 1:

		if wmiRecv.Call_id != WMI_CALL_ID_2 {
			return false
		}

		if wmiRecv.Ptype != PDU_RESP {
			return false
		}

		return true
	}

	return false
}

func NewMatcher() csm.Matcher {

	m := &WMIMatcher{}

	ds1 := DCERPC_DEFAULT{
		Rpc_ver:       5,
		Rpc_ver_minor: 0,
		Ptype:         PDU_BIND,
		Flags:         0x03,
		Drep:          0x10,
		Frag_len:      16 + 56,
		Auth_len:      0,
		Call_id:       WMI_CALL_ID_1,
	}

	ds2 := DCERPC_DEFAULT{
		Rpc_ver:       5,
		Rpc_ver_minor: 0,
		Ptype:         PDU_REQ,
		Flags:         0x03,
		Drep:          0x10,
		Frag_len:      16 + 8,
		Auth_len:      0,
		Call_id:       WMI_CALL_ID_2,
	}

	ioxidr := DCERPC_IOXIDResolver{
		MaxXmitFrag: 0x16d0,
		MaxRecvFrag: 0x16d0,

		AssocGroup: 0,
		NumCtxItem: 1,

		ContextId:    0,
		NumTransItem: 1,

		//interfaces

		InterfaceVer:      0,
		InterfaceVerMinor: 0,

		//transSyntax

		TransSyntaxVer: 2,
	}

	ioxidr.Interfaces[0] = 0xc4
	ioxidr.Interfaces[1] = 0xfe
	ioxidr.Interfaces[2] = 0xfc
	ioxidr.Interfaces[3] = 0x99

	ioxidr.Interfaces[4] = 0x60
	ioxidr.Interfaces[5] = 0x52

	ioxidr.Interfaces[6] = 0x1b
	ioxidr.Interfaces[7] = 0x10

	ioxidr.Interfaces[8] = 0xbb
	ioxidr.Interfaces[9] = 0xcb

	ioxidr.Interfaces[10] = 0x00
	ioxidr.Interfaces[11] = 0xaa
	ioxidr.Interfaces[12] = 0x00
	ioxidr.Interfaces[13] = 0x21
	ioxidr.Interfaces[14] = 0x34
	ioxidr.Interfaces[15] = 0x7a

	ioxidr.TransSyntax[0] = 0x04
	ioxidr.TransSyntax[1] = 0x5d
	ioxidr.TransSyntax[2] = 0x88
	ioxidr.TransSyntax[3] = 0x8a

	ioxidr.TransSyntax[4] = 0xeb
	ioxidr.TransSyntax[5] = 0x1c

	ioxidr.TransSyntax[6] = 0xc9
	ioxidr.TransSyntax[7] = 0x11

	ioxidr.TransSyntax[8] = 0x9f
	ioxidr.TransSyntax[9] = 0xe8

	ioxidr.TransSyntax[10] = 0x08
	ioxidr.TransSyntax[11] = 0x00
	ioxidr.TransSyntax[12] = 0x2b
	ioxidr.TransSyntax[13] = 0x10
	ioxidr.TransSyntax[14] = 0x48
	ioxidr.TransSyntax[15] = 0x60

	da := DCERPC_ALIVE{
		AllocHint: 0,
		ContextId: 0,
		OpNum:     3,
	}

	buf1 := new(bytes.Buffer)

	binary.Write(buf1, binary.LittleEndian, ds1)
	ds1Bytes := buf1.Bytes()

	buf2 := new(bytes.Buffer)

	binary.Write(buf2, binary.LittleEndian, ds2)
	ds2Bytes := buf2.Bytes()

	buf3 := new(bytes.Buffer)

	binary.Write(buf3, binary.LittleEndian, ioxidr)
	ioxidrBytes := buf3.Bytes()

	buf4 := new(bytes.Buffer)

	binary.Write(buf4, binary.LittleEndian, da)
	daBytes := buf4.Bytes()

	firstByte := make([]byte, len(ds1Bytes)+len(ioxidrBytes))

	copy(firstByte[0:], ds1Bytes)
	copy(firstByte[len(ds1Bytes):], ioxidrBytes)

	secondByte := make([]byte, len(ds2Bytes)+len(daBytes))

	copy(secondByte[0:], ds2Bytes)
	copy(secondByte[len(ds2Bytes):], daBytes)

	m.AddPacket(csm.NewPacket(firstByte, len(ds1Bytes)+len(ioxidrBytes)))
	m.AddPacket(csm.NewPacket(secondByte, len(ds2Bytes)+len(daBytes)))

	return m

}

type DCERPC_DEFAULT struct {
	Rpc_ver       uint8
	Rpc_ver_minor uint8
	Ptype         uint8
	Flags         uint8
	Drep          uint32
	Frag_len      uint16
	Auth_len      uint16
	Call_id       uint32
}

type DCERPC_ALIVE struct {
	AllocHint uint32
	ContextId uint16
	OpNum     uint16
}

type DCERPC_IOXIDResolver struct {
	MaxXmitFrag uint16
	MaxRecvFrag uint16
	AssocGroup  uint32
	NumCtxItem  uint8
	UnknownCode [3]uint8

	ContextId    uint16
	NumTransItem uint8
	UnknownCode2 uint8

	Interfaces        [16]uint8
	InterfaceVer      uint16
	InterfaceVerMinor uint16
	TransSyntax       [16]uint8
	TransSyntaxVer    uint32
}