package mongodb

import (
	"bytes"
	"encoding/binary"
	"math/rand"

	"git.loafle.net/overflow/overflow_discovery/match/packet"
	"git.loafle.net/overflow/overflow_discovery/model/scaninfo"
)

const (
	MONGO_OP_REQUEST uint32 = 2004
	MONGO_OP_REPLY   uint32 = 1
	MONGO_FCNAME     string = "admin.$cmd"
	MONGO_ELEMENT    string = "ismaster"
)

var MONGO_REQUEST_ID uint32

type mongo struct {
	MessageLength      uint32
	RequestId          uint32
	ResponseTo         uint32
	OpCode             uint32
	Flags              uint32
	FullCollectionName [11]byte
	NumberToSkip       uint32
	NumberToReturn     int32
	DocumentLength     uint32
	Type_              uint8
	Element            [9]byte
	Value              uint8
	_                  uint8
}

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

func NewMongoDBMatcher() *MongoDBMatcher {

	mongoMatcher := &MongoDBMatcher{}

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

	var fcn [11]byte
	copy(fcn[:], MONGO_FCNAME)

	var elem [9]byte
	copy(elem[:], MONGO_ELEMENT)

	MONGO_REQUEST_ID = rand.Uint32()
	m := mongo{
		MessageLength:      uint32(len(tempBuf.Bytes())),
		RequestId:          MONGO_REQUEST_ID,
		ResponseTo:         0,
		OpCode:             MONGO_OP_REQUEST,
		Flags:              0,
		FullCollectionName: fcn,
		NumberToSkip:       0,
		NumberToReturn:     -1,
		DocumentLength:     16,
		Type_:              0x08,
		Element:            elem,
		Value:              1,
	}
	writer := new(bytes.Buffer)
	binary.Write(writer, binary.LittleEndian, m)

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

	return mongoMatcher
}

func (t *MongoDBMatcher) ServiceName() string {
	return "MongoDB"
}

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

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

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

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

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

	if packet == nil {
		return false
	}

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

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

	if uint32(packet.Len) != m.MessageLength ||
		m.ResponseTo != MONGO_REQUEST_ID ||
		m.OpCode != MONGO_OP_REPLY {
		return false
	}

	return true

}