Initial Commit

socket/decoder.go

@@ -0,0 +1,8 @@
+package socket
+
+// Py_builtins_str is used by the pickle decoder to parse the server response into a format Go can understand
+type Py_builtins_str struct{}
+
+func (c Py_builtins_str) Call(args ...interface{}) (interface{}, error) {
+	return args[0], nil
+}

socket/fail2banSocket.go

@@ -0,0 +1,179 @@
+package socket
+
+import (
+	"fmt"
+	"github.com/kisielk/og-rek"
+	"github.com/nlpodyssey/gopickle/types"
+	"net"
+	"strings"
+)
+
+type Fail2BanSocket struct {
+	socket  net.Conn
+	encoder *ogórek.Encoder
+}
+
+type JailStats struct {
+	FailedCurrent int
+	FailedTotal   int
+	BannedCurrent int
+	BannedTotal   int
+}
+
+func ConnectToSocket(path string) (*Fail2BanSocket, error) {
+	c, err := net.Dial("unix", path)
+	if err != nil {
+		return nil, err
+	}
+	return &Fail2BanSocket{
+		socket:  c,
+		encoder: ogórek.NewEncoder(c),
+	}, nil
+}
+
+func (s *Fail2BanSocket) Close() error {
+	return s.socket.Close()
+}
+
+func (s *Fail2BanSocket) Ping() (bool, error) {
+	response, err := s.sendCommand([]string{pingCommand, "100"})
+	if err != nil {
+		return false, newConnectionError(pingCommand, err)
+	}
+
+	if t, ok := response.(*types.Tuple); ok {
+		if (*t)[1] == "pong" {
+			return true, nil
+		}
+		return false, fmt.Errorf("unexpected response data (expecting 'pong'): %s", (*t)[1])
+	}
+	return false, newBadFormatError(pingCommand, response)
+}
+
+func (s *Fail2BanSocket) GetJails() ([]string, error) {
+	response, err := s.sendCommand([]string{statusCommand})
+	if err != nil {
+		return nil, err
+	}
+
+	if lvl1, ok := response.(*types.Tuple); ok {
+		if lvl2, ok := lvl1.Get(1).(*types.List); ok {
+			if lvl3, ok := lvl2.Get(1).(*types.Tuple); ok {
+				if lvl4, ok := lvl3.Get(1).(string); ok {
+					splitJails := strings.Split(lvl4, ",")
+					return trimSpaceForAll(splitJails), nil
+				}
+			}
+		}
+	}
+	return nil, newBadFormatError(statusCommand, response)
+}
+
+func (s *Fail2BanSocket) GetJailStats(jail string) (JailStats, error) {
+	response, err := s.sendCommand([]string{statusCommand, jail})
+	if err != nil {
+		return JailStats{}, err
+	}
+
+	stats := JailStats{
+		FailedCurrent: -1,
+		FailedTotal:   -1,
+		BannedCurrent: -1,
+		BannedTotal:   -1,
+	}
+
+	if lvl1, ok := response.(*types.Tuple); ok {
+		if lvl2, ok := lvl1.Get(1).(*types.List); ok {
+			if filter, ok := lvl2.Get(0).(*types.Tuple); ok {
+				if filterLvl1, ok := filter.Get(1).(*types.List); ok {
+					if filterCurrentTuple, ok := filterLvl1.Get(0).(*types.Tuple); ok {
+						if filterCurrent, ok := filterCurrentTuple.Get(1).(int); ok {
+							stats.FailedCurrent = filterCurrent
+						}
+					}
+					if filterTotalTuple, ok := filterLvl1.Get(1).(*types.Tuple); ok {
+						if filterTotal, ok := filterTotalTuple.Get(1).(int); ok {
+							stats.FailedTotal = filterTotal
+						}
+					}
+				}
+			}
+			if actions, ok := lvl2.Get(1).(*types.Tuple); ok {
+				if actionsLvl1, ok := actions.Get(1).(*types.List); ok {
+					if actionsCurrentTuple, ok := actionsLvl1.Get(0).(*types.Tuple); ok {
+						if actionsCurrent, ok := actionsCurrentTuple.Get(1).(int); ok {
+							stats.BannedCurrent = actionsCurrent
+						}
+					}
+					if actionsTotalTuple, ok := actionsLvl1.Get(1).(*types.Tuple); ok {
+						if actionsTotal, ok := actionsTotalTuple.Get(1).(int); ok {
+							stats.BannedTotal = actionsTotal
+						}
+					}
+				}
+			}
+			return stats, nil
+		}
+	}
+	return stats, newBadFormatError(statusCommand, response)
+}
+
+func (s *Fail2BanSocket) GetJailBanTime(jail string) (int, error) {
+	command := fmt.Sprintf(banTimeCommandFmt, jail)
+	return s.sendSimpleIntCommand(command)
+}
+
+func (s *Fail2BanSocket) GetJailFindTime(jail string) (int, error) {
+	command := fmt.Sprintf(findTimeCommandFmt, jail)
+	return s.sendSimpleIntCommand(command)
+}
+
+func (s *Fail2BanSocket) GetJailMaxRetries(jail string) (int, error) {
+	command := fmt.Sprintf(maxRetriesCommandFmt, jail)
+	return s.sendSimpleIntCommand(command)
+}
+
+func (s *Fail2BanSocket) GetServerVersion() (string, error) {
+	response, err := s.sendCommand([]string{versionCommand})
+	if err != nil {
+		return "", err
+	}
+
+	if lvl1, ok := response.(*types.Tuple); ok {
+		if versionStr, ok := lvl1.Get(1).(string); ok {
+			return versionStr, nil
+		}
+	}
+	return "", newBadFormatError(versionCommand, response)
+}
+
+// sendSimpleIntCommand sends a command to the fail2ban socket and parses the response to extract an int.
+// This command assumes that the response data is in the format of `(d, d)` where `d` is a number.
+func (s *Fail2BanSocket) sendSimpleIntCommand(command string) (int, error) {
+	response, err := s.sendCommand(strings.Split(command, " "))
+	if err != nil {
+		return -1, err
+	}
+
+	if lvl1, ok := response.(*types.Tuple); ok {
+		if banTime, ok := lvl1.Get(1).(int); ok {
+			return banTime, nil
+		}
+	}
+	return -1, newBadFormatError(command, response)
+}
+
+func newBadFormatError(command string, data interface{}) error {
+	return fmt.Errorf("(%s) unexpected response format - cannot parse: %v", command, data)
+}
+
+func newConnectionError(command string, err error) error {
+	return fmt.Errorf("(%s) failed to send command through socket: %v", command, err)
+}
+
+func trimSpaceForAll(slice []string) []string {
+	for i := range slice {
+		slice[i] = strings.TrimSpace(slice[i])
+	}
+	return slice
+}

socket/protocol.go

@@ -0,0 +1,69 @@
+package socket
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"github.com/nlpodyssey/gopickle/pickle"
+)
+
+const (
+	commandTerminator    = "<F2B_END_COMMAND>"
+	pingCommand          = "ping"
+	statusCommand        = "status"
+	versionCommand       = "version"
+	banTimeCommandFmt    = "get %s bantime"
+	findTimeCommandFmt   = "get %s findtime"
+	maxRetriesCommandFmt = "get %s maxretry"
+	socketReadBufferSize = 1024
+)
+
+func (s *Fail2BanSocket) sendCommand(command []string) (interface{}, error) {
+	err := s.write(command)
+	if err != nil {
+		return nil, err
+	}
+	return s.read()
+}
+
+func (s *Fail2BanSocket) write(command []string) error {
+	err := s.encoder.Encode(command)
+	if err != nil {
+		return err
+	}
+	_, err = s.socket.Write([]byte(commandTerminator))
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (s *Fail2BanSocket) read() (interface{}, error) {
+	reader := bufio.NewReader(s.socket)
+
+	data := []byte{}
+	for {
+		buf := make([]byte, socketReadBufferSize)
+		_, err := reader.Read(buf)
+		if err != nil {
+			return nil, err
+		}
+		data = append(data, buf...)
+		containsTerminator := bytes.Contains(data, []byte(commandTerminator))
+		if containsTerminator {
+			break
+		}
+	}
+
+	bufReader := bytes.NewReader(data)
+	unpickler := pickle.NewUnpickler(bufReader)
+
+	unpickler.FindClass = func(module, name string) (interface{}, error) {
+		if (module == "builtins" || module == "__builtin__") && name == "str" {
+			return &Py_builtins_str{}, nil
+		}
+		return nil, fmt.Errorf("class not found: " + module + " : " + name)
+	}
+
+	return unpickler.Load()
+}