package services

import (
    "net"
    "regexp"
    "strings"
    "time"
)

func NewPostgreSQLHandler(log LoggerFunc) Handler {
	return func(conn net.Conn) {
		defer conn.Close()
		remote := conn.RemoteAddr().String()
		conn.SetDeadline(time.Now().Add(30 * time.Second))
		buf := make([]byte, 1024)
		n, err := conn.Read(buf)
		if err != nil { return }
		if n >= 8 {
			length := int(buf[0])<<24 | int(buf[1])<<16 | int(buf[2])<<8 | int(buf[3])
			if length > 8 && n >= length {
				params := string(buf[8:length])
				if strings.Contains(params, "user") {
					re := regexp.MustCompile(`user\x00([^\x00]+)`) 
					m := re.FindStringSubmatch(params)
					if len(m) > 1 {
						log(Record{Timestamp: Now(), RemoteAddr: remoteIP(remote), RemotePort: remotePort(remote), Service: "postgresql", Details: map[string]string{"event":"auth_attempt","username":m[1]}})
					}
				}
			}
		}
		// request cleartext password
		authReq := []byte{0x52, 0x00,0x00,0x00,0x08, 0x00,0x00,0x00,0x03}
		_, _ = conn.Write(authReq)
		n, err = conn.Read(buf)
		if err == nil && n > 5 && buf[0] == 0x70 {
			length := int(buf[1])<<24 | int(buf[2])<<16 | int(buf[3])<<8 | int(buf[4])
			if length > 5 && n >= length {
				password := string(buf[5:length-1])
				log(Record{Timestamp: Now(), RemoteAddr: remoteIP(remote), RemotePort: remotePort(remote), Service: "postgresql", Details: map[string]string{"event":"password_attempt","password":password}})
			}
		}
		// error
		errMsg := []byte{0x45, 0x00,0x00,0x00,0x26}
		errMsg = append(errMsg, []byte("SFATAL\x00C28P01\x00Mpassword authentication failed\x00\x00")...)
		_, _ = conn.Write(errMsg)
	}
}