perf: 优化Log.go的代码,添加注释,规范输出

This commit is contained in:
ZacharyZcR
2024-12-18 22:24:11 +08:00
parent 9433741471
commit 1d0676e508
+65 -38
View File
@@ -12,49 +12,60 @@ import (
"time" "time"
) )
var Num int64 // 记录扫描状态的全局变量
var End int64 var (
var Results = make(chan *string) Num int64 // 总任务数
var LogSucTime int64 End int64 // 已完成数
var LogErrTime int64 Results = make(chan *string) // 结果通道
var WaitTime int64 LogSucTime int64 // 最近成功日志时间
var Silent bool LogErrTime int64 // 最近错误日志时间
var Nocolor bool WaitTime int64 // 等待时间
var JsonOutput bool Silent bool // 静默模式
var LogWG sync.WaitGroup Nocolor bool // 禁用颜色
JsonOutput bool // JSON输出
LogWG sync.WaitGroup // 日志同步等待组
)
// JsonText JSON输出的结构体
type JsonText struct { type JsonText struct {
Type string `json:"type"` Type string `json:"type"` // 消息类型
Text string `json:"text"` Text string `json:"text"` // 消息内容
} }
// init 初始化日志配置
func init() { func init() {
log.SetOutput(io.Discard) log.SetOutput(io.Discard)
LogSucTime = time.Now().Unix() LogSucTime = time.Now().Unix()
go SaveLog() go SaveLog()
} }
// LogSuccess 记录成功信息
func LogSuccess(result string) { func LogSuccess(result string) {
LogWG.Add(1) LogWG.Add(1)
LogSucTime = time.Now().Unix() LogSucTime = time.Now().Unix()
Results <- &result Results <- &result
} }
// SaveLog 保存日志信息
func SaveLog() { func SaveLog() {
for result := range Results { for result := range Results {
// 打印日志
if !Silent { if !Silent {
if Nocolor { if Nocolor {
fmt.Println(*result) fmt.Println(*result)
} else { } else {
if strings.HasPrefix(*result, "[+] InfoScan") { switch {
case strings.HasPrefix(*result, "[+] 信息扫描"):
color.Green(*result) color.Green(*result)
} else if strings.HasPrefix(*result, "[+]") { case strings.HasPrefix(*result, "[+]"):
color.Red(*result) color.Red(*result)
} else { default:
fmt.Println(*result) fmt.Println(*result)
} }
} }
} }
// 保存到文件
if IsSave { if IsSave {
WriteFile(*result, Outputfile) WriteFile(*result, Outputfile)
} }
@@ -62,17 +73,20 @@ func SaveLog() {
} }
} }
// WriteFile 写入文件
func WriteFile(result string, filename string) { func WriteFile(result string, filename string) {
// 打开文件
fl, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666) fl, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
if err != nil { if err != nil {
fmt.Printf("Open %s error, %v\n", filename, err) fmt.Printf("[!] 打开文件失败 %s: %v\n", filename, err)
return return
} }
defer fl.Close()
if JsonOutput { if JsonOutput {
var scantype string // 解析JSON格式
var text string var scantype, text string
if strings.HasPrefix(result, "[+]") || strings.HasPrefix(result, "[*]") || strings.HasPrefix(result, "[-]") { if strings.HasPrefix(result, "[+]") || strings.HasPrefix(result, "[*]") || strings.HasPrefix(result, "[-]") {
//找到第二个空格的位置
index := strings.Index(result[4:], " ") index := strings.Index(result[4:], " ")
if index == -1 { if index == -1 {
scantype = "msg" scantype = "msg"
@@ -85,61 +99,74 @@ func WriteFile(result string, filename string) {
scantype = "msg" scantype = "msg"
text = result text = result
} }
// 构造JSON对象
jsonText := JsonText{ jsonText := JsonText{
Type: scantype, Type: scantype,
Text: text, Text: text,
} }
// 序列化JSON
jsonData, err := json.Marshal(jsonText) jsonData, err := json.Marshal(jsonText)
if err != nil { if err != nil {
fmt.Println(err) fmt.Printf("[!] JSON序列化失败: %v\n", err)
jsonText = JsonText{ jsonText = JsonText{
Type: "msg", Type: "msg",
Text: result, Text: result,
} }
jsonData, err = json.Marshal(jsonText) jsonData, _ = json.Marshal(jsonText)
if err != nil {
fmt.Println(err)
jsonData = []byte(result)
}
} }
jsonData = append(jsonData, []byte(",\n")...) jsonData = append(jsonData, []byte(",\n")...)
_, err = fl.Write(jsonData) _, err = fl.Write(jsonData)
} else { } else {
_, err = fl.Write([]byte(result + "\n")) _, err = fl.Write([]byte(result + "\n"))
} }
fl.Close()
if err != nil { if err != nil {
fmt.Printf("Write %s error, %v\n", filename, err) fmt.Printf("[!] 写入文件失败 %s: %v\n", filename, err)
} }
} }
// LogError 记录错误信息
func LogError(errinfo interface{}) { func LogError(errinfo interface{}) {
if WaitTime == 0 { if WaitTime == 0 {
fmt.Printf("已完成 %v/%v %v \n", End, Num, errinfo) fmt.Printf("[*] 已完成 %v/%v %v\n", End, Num, errinfo)
} else if (time.Now().Unix()-LogSucTime) > WaitTime && (time.Now().Unix()-LogErrTime) > WaitTime { } else if (time.Now().Unix()-LogSucTime) > WaitTime && (time.Now().Unix()-LogErrTime) > WaitTime {
fmt.Printf("已完成 %v/%v %v \n", End, Num, errinfo) fmt.Printf("[*] 已完成 %v/%v %v\n", End, Num, errinfo)
LogErrTime = time.Now().Unix() LogErrTime = time.Now().Unix()
} }
} }
// CheckErrs 检查是否为已知错误
func CheckErrs(err error) bool { func CheckErrs(err error) bool {
if err == nil { if err == nil {
return false return false
} }
// 已知错误列表
errs := []string{ errs := []string{
"closed by the remote host", "too many connections", "远程主机关闭连接",
"i/o timeout", "EOF", "A connection attempt failed", "连接数过多",
"established connection failed", "connection attempt failed", "连接超时",
"Unable to read", "is not allowed to connect to this", "EOF",
"no pg_hba.conf entry", "连接尝试失败",
"No connection could be made", "建立连接失败",
"invalid packet size", "无法连接",
"bad connection", "无法读取数据",
"不允许连接",
"无pg_hba.conf配置项",
"无法建立连接",
"无效的数据包大小",
"连接异常",
} }
// 检查错误是否匹配
errLower := strings.ToLower(err.Error())
for _, key := range errs { for _, key := range errs {
if strings.Contains(strings.ToLower(err.Error()), strings.ToLower(key)) { if strings.Contains(errLower, strings.ToLower(key)) {
return true return true
} }
} }
return false return false
} }