network_service_matcher/postgresql/postgresql.go

package postgresql

import (
	"bytes"
	"encoding/binary"
	"strings"

	cnsm "git.loafle.net/commons_go/network_service_matcher"
)

const (
	RESPONSE_TYPE_ERR uint8 = 0x45
)

type pgsql struct {
	Len              uint32
	MessageType      uint16
	_                uint16
	Name             [5]byte
	NameValue        byte
	Db               [9]byte
	DBValue          byte
	Encoding         [16]byte
	EncodingValue    [5]byte
	DateStyle        [10]byte
	DateStyleValue   [4]byte
	TimeZone         [9]byte
	TimeZoneValue    [11]byte
	ExtraDigits      [19]byte
	ExtraDigitsValue uint16
	End              byte
}

type pgsqlErrResponse struct {
	ResponseType uint8
	Len          [4]byte
	Severity     [6]byte
	_            byte
	Code         [6]byte
	_            byte
	Message      [53]byte
}

type pgsqlErrResponse2 struct {
	ResponseType uint8
	Len          [4]byte
	Data         [128]byte
	//Severity     [6]byte
	//_            byte
	//Code         [6]byte
	//_            byte
	//Message      [53]byte
}

type PostgreSQLMatcher struct {
	cnsm.Matchers
}

func (t *PostgreSQLMatcher) ServiceName() string {
	return "PostgreSQL"
}

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

func (t *PostgreSQLMatcher) HasResponse(index int) bool {
	return true
}

func (t *PostgreSQLMatcher) IsError(info cnsm.MatchInfo, index int, packet *cnsm.Packet) bool {
	return false
}

func (t *PostgreSQLMatcher) Match(info cnsm.MatchInfo, index int, packet *cnsm.Packet) bool {

	if packet == nil {
		return false
	}

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

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

	if pg.ResponseType != RESPONSE_TYPE_ERR {
		return false
	}

	length := binary.BigEndian.Uint32(pg.Len[:])
	if length+1 != uint32(packet.Len) {
		return false
	}

	data := string(pg.Data[:])
	splits := strings.Split(data, "\x00")

	var findSeverity bool = false
	var findErrorCode bool = false
	for _, s := range splits {
		if strings.Contains(s, "FATAL") {
			findSeverity = true
		}
		if strings.Contains(s, "28000") {
			findErrorCode = true
		}
	}

	if !findSeverity || !findErrorCode {
		return false
	}

	return true

}

func NewMatcher() cnsm.Matcher {

	m := &PostgreSQLMatcher{}

	pg := pgsql{}

	pg.Len = 0x00000065
	pg.MessageType = 0x0003

	var name [5]byte
	copy(name[:], "user")
	pg.Name = name
	pg.NameValue = 0x00

	var db [9]byte
	copy(db[:], "database")
	pg.Db = db
	pg.DBValue = 0x00

	var encoding [16]byte
	copy(encoding[:], "client_encoding")
	pg.Encoding = encoding

	var encodingValue [5]byte
	copy(encodingValue[:], "UTF8")
	pg.EncodingValue = encodingValue

	var dateStyle [10]byte
	copy(dateStyle[:], "DateStyle")
	pg.DateStyle = dateStyle

	var dateStyleValue [4]byte
	copy(dateStyleValue[:], "ISO")
	pg.DateStyleValue = dateStyleValue

	var timeZone [9]byte
	copy(timeZone[:], "TimeZone")
	pg.TimeZone = timeZone

	var timeZoneValue [11]byte
	copy(timeZoneValue[:], "Asia/Seoul")
	pg.TimeZoneValue = timeZoneValue

	var extraDigit [19]byte
	copy(extraDigit[:], "extra_float_digits")
	pg.ExtraDigits = extraDigit

	pg.ExtraDigitsValue = 0x3200

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

	m.AddPacket(cnsm.NewPacket(writer.Bytes(), writer.Len()))

	return m
}