package mssql

import (
	"bytes"
	"encoding/binary"
	"git.loafle.net/overflow/overflow_probe/matcher/packet"
	"git.loafle.net/overflow/overflow_probe/model/scaninfo"
)

const (
	HEADER_TYPE_PRELOGIN       uint8 = 0x12
	HEADER_TYPE_RESPONSE       uint8 = 0x4
	PL_OPTION_TOKEN_VERSION    uint8 = 0x00
	PL_OPTION_TOKEN_ENCRYPTION uint8 = 0x01
	PL_OPTION_TOKEN_TRACEID    uint8 = 0x05
	PL_OPTION_TOKEN_TERMINATOR uint8 = 0xff

	ENCRYPT_OFF     string = "Encryption is available but off."
	ENCRYPT_ON      string = "Encryption is available and on."
	ENCRYPT_NOT_SUP string = "Encryption is not available."
	ENCRYPT_REQ     string = "Encryption is required."
)

type PreloginMsg struct {
	VersionToken  uint8
	VersionOffset uint16
	VersionLength uint16

	EncryptionToken  uint8
	EncryptionOffset uint16
	EncryptionLength uint16

	TraceIdToken  uint8
	TraceIdOffset uint16
	TraceIdLength uint16

	Terminator uint8
	Options    [7]uint8
	TraceId    [36]uint8
}

type mssql struct {
	Type_     uint8
	Status    uint8
	Length    uint16
	Channel   uint16
	PacketNum uint8
	Window    uint8
	Prelogin  PreloginMsg
}

type PreloginResponse struct {
	Msg [256]uint8
}

type mssqlResponse struct {
	Type_        uint8
	Status       uint8
	Length       uint16
	Channel      uint16
	PacketNum    uint8
	Window       uint8
	PreLoginResp PreloginResponse
}

type MSSqlMatcher struct {
	packets []*packet.Packet
	isSSL   bool
}

func NewMSSqlMatcher() *MSSqlMatcher {

	mssqlMatcher := &MSSqlMatcher{}

	tempBuf := new(bytes.Buffer)
	binary.Write(tempBuf, binary.BigEndian, mssql{})

	m := mssql{
		Type_:     HEADER_TYPE_PRELOGIN,
		Status:    0x01,
		Length:    uint16(len(tempBuf.Bytes())),
		Channel:   0,
		PacketNum: 0,
		Window:    0,
		Prelogin: PreloginMsg{
			VersionToken:     PL_OPTION_TOKEN_VERSION,
			VersionOffset:    0x0010,
			VersionLength:    0x0006,
			EncryptionToken:  PL_OPTION_TOKEN_ENCRYPTION,
			EncryptionOffset: 0x0016,
			EncryptionLength: 0x0001,
			TraceIdToken:     PL_OPTION_TOKEN_TRACEID,
			TraceIdOffset:    0x0017,
			TraceIdLength:    0x0024,
			Terminator:       PL_OPTION_TOKEN_TERMINATOR,
		},
	}

	writer := new(bytes.Buffer)
	binary.Write(writer, binary.BigEndian, m)

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

	return mssqlMatcher
}

func (t *MSSqlMatcher) ServiceName() string {
	if t.isSSL {
		return "SQL Server (SSL)"
	}
	return "SQL Server"
}

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

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

func (t *MSSqlMatcher) IsNoResponse(index int) bool {
	return false
}

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

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

	if packet == nil {
		return false
	}

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

	m := mssqlResponse{}

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

	if m.Type_ != HEADER_TYPE_RESPONSE {
		return false
	}

	if m.Length != uint16(packet.Len) {
		return false
	}

	switch m.PreLoginResp.Msg[m.Length-9 : m.Length-8][0] {
	case 0:
		return true
	case 1:
		t.isSSL = true
		return true
	case 2:
		return true
	case 3:
		t.isSSL = true
		return true
	default:
		return false
	}

	return false

}
