[Linux]ban sshd fail Login IP with Generate ban list

Ban sshd fail login ip.

package util

import (
	"fmt"
	"github.com/stretchr/testify/assert"
	"io/ioutil"
	"os"
	"sort"
	"strings"
	"testing"
	"time"
)

func TestRead(t *testing.T) {
	path := "C:\\Users\\v_zhenlluo\\Downloads\\sshdLog"
	fileList, e := ioutil.ReadDir(path)
	if e != nil {
		fmt.Println("read dir error")
		return
	}
	ipAccessCountMap := make(map[string]int) // sshd访问失败 统计次数
	// 提取IP
	for _, fileInfo := range fileList {
		if !strings.HasPrefix(fileInfo.Name(), "secure") {
			continue
		}
		filePath := fmt.Sprintf("%s\\%s", path, fileInfo.Name())
		content, err := os.ReadFile(filePath)
		if err != nil {
			continue
		}
		splitStrItems := strings.Split(string(content), "\n")
		for _, item := range splitStrItems {
			sshdFailLoginLog := strings.ToLower(item)
			if !(strings.Contains(sshdFailLoginLog, "invalid") && strings.Contains(sshdFailLoginLog, "sshd")) {
				continue
			}
			sp1 := strings.Split(sshdFailLoginLog, "from")
			if len(sp1) < 2 {
				continue
			}
			sp2 := strings.Split(sp1[1], " ")
			if len(sp2) < 2 {
				continue
			}
			accessIP := sp2[1]
			sp3 := strings.Split(accessIP, ".")
			if len(sp3) < 4 {
				continue
			}
			ipAccessCountMap[accessIP]++
		}
	}
	// 排序
	var ipAccessList []IpAccessInfo
	for accessIP, count := range ipAccessCountMap {
		ipAccessList = append(ipAccessList, IpAccessInfo{Ip: accessIP, Count: count})
	}
	sort.Sort(IpAccess(ipAccessList))
	// 登陆失败10次 封禁
	banIpBlockMap := make(map[string]*HashSetString) // 最终封禁IP列表
	for _, info := range ipAccessList {
		if info.Count < 10 {
			continue
		}
		accessIP := info.Ip
		sp3 := strings.Split(accessIP, ".")
		if len(sp3) < 4 {
			continue
		}
		banIpBlock := fmt.Sprintf("%s.%s.0.0/24", sp3[0], sp3[1])
		if _, ok := banIpBlockMap[banIpBlock]; !ok {
			banIpBlockMap[banIpBlock] = NewHashSetString()
		}
		banIpBlockMap[banIpBlock].Add(accessIP)
	}
	// 初始化文件
	writeFilePath := fmt.Sprintf("%s\\ban_%s.txt", path, time.Now().Format("20060102150405"))
	f, err := os.OpenFile(writeFilePath, os.O_CREATE, 0777)
	defer f.Close()
	if err != nil {
		return
	}
	// 写入
	for banIp, _ := range banIpBlockMap {
		f.Write([]byte(banIp)) //data是自己定义的数据
		f.WriteString("\n")
	}
	assert.Equal(t, len(ipAccessCountMap), len(ipAccessCountMap))
}

object order:

type IpAccessInfo struct {
	Ip    string
	Count int
}

type IpAccess []IpAccessInfo

func (a IpAccess) Len() int           { return len(a) }
func (a IpAccess) Less(i, j int) bool { return a[i].Count > a[j].Count } // desc order
func (a IpAccess) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

hashSetString Util:

// HashSetString 字符串hash结构
type HashSetString struct {
	dic  map[string]interface{} // 字典
	Data []string               // 数据
}

// NewHashSetString 初始化
func NewHashSetString() *HashSetString {
	return &HashSetString{
		dic:  make(map[string]interface{}, 0),
		Data: make([]string, 0),
	}
}

// NewHashSetStringWithData 初始化
func NewHashSetStringWithData(data []string) *HashSetString {
	set := &HashSetString{
		dic:  make(map[string]interface{}, 0),
		Data: make([]string, 0),
	}
	for i := range data {
		set.Add(data[i])
	}
	return set
}

// Add 新增元素
func (h *HashSetString) Add(ele string) bool {
	_, ok := h.dic[ele]
	if ok {
		return false
	}
	h.dic[ele] = ele
	h.Data = append(h.Data, ele)
	return true
}

// Contains 是否存在字符串
func (h *HashSetString) Contains(ele string) bool {
	_, ok := h.dic[ele]
	if ok {
		return true
	}
	return false
}

// Get 获取元素
func (h *HashSetString) Get(ele string) string {
	val, ok := h.dic[ele]
	if ok {
		return val.(string)
	}
	return ""
}

// GetList 获取所有字符串
func (h *HashSetString) GetList() []string {
	return h.Data
}

// HashSetInt 数字hash结构
type HashSetInt struct {
	dic  map[int]interface{}
	Data []int
}

// NewHashSetInt 初始化
func NewHashSetInt() *HashSetInt {
	return &HashSetInt{
		dic:  make(map[int]interface{}, 0),
		Data: make([]int, 0),
	}
}

// Add 新增元素
func (h *HashSetInt) Add(ele int) bool {
	_, ok := h.dic[ele]
	if ok {
		return false
	}
	h.dic[ele] = ele
	h.Data = append(h.Data, ele)
	return true
}

// Get 获取元素
func (h *HashSetInt) Get(ele int) int {
	val, ok := h.dic[ele]
	if ok {
		return val.(int)
	}
	return -1
}

// GetList 获取所有数字
func (h *HashSetInt) GetList() []int {
	return h.Data
}