jar patcher is able to remove JndiLookup.class file from jars

tools/log4shell/analyze/analyze.go

@@ -97,12 +97,10 @@ func fileNameToSemver(fileNameNoExt string) string {
 	return semverVersion
 }
 
-func GetJndiLookupHash(zipReader *zip.Reader, filePath string) (fileName, fileHash string, err error) {
-	fileName = "org/apache/logging/log4j/core/lookup/JndiLookup.class"
-
-	reader, err := zipReader.Open(fileName)
+func GetJndiLookupHash(zipReader *zip.Reader, filePath string) (fileName, fileHash string) {
+	reader, err := zipReader.Open(constants.JndiLookupClasspath)
 	if err != nil {
-		log.Warn().
+		log.Debug().
 			Str("fieName", fileName).
 			Str("path", filePath).
 			Err(err).
@@ -113,7 +111,7 @@ func GetJndiLookupHash(zipReader *zip.Reader, filePath string) (fileName, fileHa
 
 	fileHash, err = util.HexEncodedSha256FromReader(reader)
 	if err != nil {
-		log.Warn().
+		log.Debug().
 			Str("fieName", fileName).
 			Str("path", filePath).
 			Err(err).
@@ -124,6 +122,11 @@ func GetJndiLookupHash(zipReader *zip.Reader, filePath string) (fileName, fileHa
 }
 
 func ProcessArchiveFile(zipReader *zip.Reader, reader io.Reader, filePath, fileName string) (finding *types.Finding) {
+	var (
+		jndiLookupFileName string
+		jndiLookupFileHash string
+	)
+
 	_, file := path.Split(filePath)
 	fileNameNoExt := strings.TrimSuffix(file, path.Ext(file))
 
@@ -163,15 +166,8 @@ func ProcessArchiveFile(zipReader *zip.Reader, reader io.Reader, filePath, fileN
 		return
 	}
 
-	jndiLookupFileName := ""
-	jndiLookupFileHash := ""
-
 	if versionIsInRange(fileNameNoExt, semverVersion, constants.JndiLookupPatchFileVersions) {
-		jndiLookupFileName, jndiLookupFileHash, err = GetJndiLookupHash(zipReader, filePath)
-		if err != nil {
-			jndiLookupFileName = ""
-			jndiLookupFileHash = ""
-		}
+		jndiLookupFileName, jndiLookupFileHash = GetJndiLookupHash(zipReader, filePath)
 	}
 
 	log.Log().

tools/log4shell/commands/patch.go

@@ -17,109 +17,270 @@ package commands
 import (
 	"archive/zip"
 	"encoding/json"
+	"fmt"
+	"github.com/lunasec-io/lunasec/tools/log4shell/scan"
 	"github.com/lunasec-io/lunasec/tools/log4shell/types"
 	"github.com/lunasec-io/lunasec/tools/log4shell/util"
 	"github.com/rs/zerolog/log"
 	"github.com/urfave/cli/v2"
-	"io/fs"
 	"io/ioutil"
 	"os"
 )
 
-func JavaArchivePatchCommand(c *cli.Context, globalBoolFlags map[string]bool) error {
-	enableGlobalFlags(c, globalBoolFlags)
-
-	findingsFile := c.String("findings")
+func scanForFindings(
+	log4jLibraryHashes []byte,
+	searchDirs []string,
+	excludeDirs []string,
+	noFollowSymlinks bool,
+) (findings []types.Finding, err error) {
+	var (
+		hashLookup types.VulnerableHashLookup
+	)
 
-	findingsContent, err := ioutil.ReadFile(findingsFile)
+	hashLookup, err = loadHashLookup(log4jLibraryHashes, "", false)
 	if err != nil {
-		log.Error().
-			Err(err).
-			Str("findings", findingsFile).
-			Msg("Unable to open and read findings file")
-		return err
+		return
 	}
 
-	var findings types.FindingsOutput
-	err = json.Unmarshal(findingsContent, &findings)
-	if err != nil {
-		log.Error().
-			Err(err).
-			Str("findings", findingsFile).
-			Msg("Unable to unmarshal findings file")
-		return err
-	}
+	processArchiveFile := scan.IdentifyPotentiallyVulnerableFiles(false, hashLookup)
+
+	scanner := scan.NewLog4jDirectoryScanner(
+		excludeDirs, false, noFollowSymlinks, processArchiveFile)
+
+	findings = scanner.Scan(searchDirs)
+	return
+}
 
-	for _, finding := range findings.VulnerableLibraries {
-		var file *os.File
+func loadOrScanForFindings(
+	c *cli.Context,
+	log4jLibraryHashes []byte,
+) (findings []types.Finding, err error) {
+	findingsFile := c.String("findings")
+	if findingsFile != "" {
+		var (
+			findingsContent []byte
+			findingsOutput types.FindingsOutput
+		)
 
-		file, err = os.Open(finding.Path)
+		findingsContent, err = ioutil.ReadFile(findingsFile)
 		if err != nil {
-			log.Warn().
-				Str("path", finding.Path).
+			log.Error().
 				Err(err).
-				Msg("Unable to open findings archive")
-			return err
+				Str("findings", findingsFile).
+				Msg("Unable to open and read findings file")
+			return
 		}
-		defer file.Close()
 
-		info, _ := os.Stat(finding.Path)
+		err = json.Unmarshal(findingsContent, &findingsOutput)
+		if err != nil {
+			log.Error().
+				Err(err).
+				Str("findings", findingsFile).
+				Msg("Unable to unmarshal findings file")
+			return
+		}
+		findings = findingsOutput.VulnerableLibraries
+		return
+	}
+
+	searchDirs := c.Args().Slice()
+
+	excludeDirs := c.StringSlice("exclude")
+	noFollowSymlinks := c.Bool("no-follow-symlinks")
+
+	log.Info().
+		Strs("searchDirs", searchDirs).
+		Strs("excludeDirs", excludeDirs).
+		Msg("Scanning directories for vulnerable Log4j libraries.")
+
+	return scanForFindings(log4jLibraryHashes, searchDirs, excludeDirs, noFollowSymlinks)
+}
 
-		var zipReader *zip.Reader
+func askIfShouldSkipLibrary(msg string) (shouldSkip, forcePatch bool) {
+	var (
+		patchPromptResp string
+	)
 
-		zipReader, err = zip.NewReader(file, info.Size())
+	for {
+		fmt.Printf("Are you sure you want to patch: %s? (y)es/(n)o/(a)ll: ", msg)
+		_, err := fmt.Scan(&patchPromptResp)
 		if err != nil {
-			log.Warn().
-				Str("path", finding.Path).
+			log.Error().
 				Err(err).
-				Msg("Unable to open archive for patching")
-			return err
+				Msg("Unable to process response.")
+			return true, false
 		}
+		fmt.Println()
 
-		var zipFile fs.File
+		switch patchPromptResp {
+		case "y":
+			shouldSkip = false
+		case "n":
+			shouldSkip = true
+		case "a":
+			forcePatch = true
+		default:
+			fmt.Printf("Option %s is not valid, please enter 'y', 'n', or 'a'.\n", patchPromptResp)
+			continue
+		}
+		break
+	}
+	return
+}
 
-		if finding.JndiLookupFileName == "" {
-			log.Warn().
+func filterOutJndiLookupFromZip(
+	finding types.Finding,
+	zipReader *zip.Reader,
+	writer *zip.Writer,
+) error {
+	for _, member := range zipReader.File {
+		if member.Name == finding.JndiLookupFileName {
+			log.Debug().
 				Str("path", finding.Path).
-				Err(err).
-				Msg("Finding does not have JndiLookup.class file to patch")
+				Str("zipFilePath", finding.JndiLookupFileName).
+				Msg("Found file to remove in order to patch log4j library.")
 			continue
 		}
 
-		zipFile, err = zipReader.Open(finding.JndiLookupFileName)
-		if err != nil {
-			log.Warn().
-				Str("path", finding.Path).
-				Str("jndiLookupFileName", finding.JndiLookupFileName).
+		if err := writer.Copy(member); err != nil {
+			log.Error().
 				Err(err).
-				Msg("Unable to open file from zip")
+				Msg("Error while copying zip file.")
 			return err
 		}
+	}
+	return nil
+}
 
-		var zipFileHash string
+func patchJavaArchive(finding types.Finding) (err error) {
+	var (
+		libraryFile *os.File
+		zipReader   *zip.Reader
+	)
 
-		zipFileHash, err = util.HexEncodedSha256FromReader(zipFile)
-		if err != nil {
+	libraryFile, err = os.Open(finding.Path)
+	if err != nil {
+		log.Error().
+			Str("path", finding.Path).
+			Err(err).
+			Msg("Unable to open findings archive")
+		return
+	}
+	defer libraryFile.Close()
+
+	info, _ := os.Stat(finding.Path)
+
+	zipReader, err = zip.NewReader(libraryFile, info.Size())
+	if err != nil {
+		log.Error().
+			Str("path", finding.Path).
+			Err(err).
+			Msg("Unable to open archive for patching")
+		return
+	}
+
+	outZip, err := ioutil.TempFile(os.TempDir(), "*.zip")
+	if err != nil {
+		log.Error().
+			Str("tmpDir", os.TempDir()).
+			Err(err).
+			Msg("Unable to create temporary libraryFile")
+		return
+	}
+	defer os.Remove(outZip.Name())
+
+	writer := zip.NewWriter(outZip)
+	defer writer.Close()
+
+	err = filterOutJndiLookupFromZip(finding, zipReader, writer)
+	if err != nil {
+		return
+	}
+
+	writer.Close()
+
+	if err = libraryFile.Close(); err != nil {
+		log.Error().
+			Str("outZipName", outZip.Name()).
+			Str("libraryFileName", finding.Path).
+			Err(err).
+			Msg("Unable to close library file.")
+		return
+	}
+
+	if err = outZip.Close(); err != nil {
+		log.Error().
+			Str("outZipName", outZip.Name()).
+			Str("libraryFileName", finding.Path).
+			Err(err).
+			Msg("Unable to close output zip.")
+		return
+	}
+
+	_, err = util.CopyFile(outZip.Name(), finding.Path)
+	if err != nil {
+		log.Error().
+			Str("outZipName", outZip.Name()).
+			Str("libraryFileName", finding.Path).
+			Err(err).
+			Msg("Unable to replace library file with patched library file.")
+		return
+	}
+	return
+}
+
+func JavaArchivePatchCommand(
+	c *cli.Context,
+	globalBoolFlags map[string]bool,
+	log4jLibraryHashes []byte,
+) error {
+	enableGlobalFlags(c, globalBoolFlags)
+
+	findings, err := loadOrScanForFindings(c, log4jLibraryHashes)
+	if err != nil {
+		return err
+	}
+
+	log.Info().
+		Int("findingsCount", len(findings)).
+		Msg("Patching found vulnerable Log4j libraries.")
+
+	forcePatch := c.Bool("force-patch")
+
+	var patchedLibraries []string
+
+	for _, finding := range findings {
+		var (
+			shouldSkip bool
+		)
+
+		if finding.JndiLookupFileName == "" {
 			log.Warn().
 				Str("path", finding.Path).
 				Err(err).
-				Msg("Unable to hash zip file")
-			return err
+				Msg("Finding does not have JndiLookup.class file to patch")
+			continue
 		}
 
-		if zipFileHash != finding.JndiLookupHash {
-			log.Warn().
-				Str("path", finding.Path).
-				Str("hash", finding.JndiLookupHash).
-				Err(err).
-				Msg("Hashes do not match, not deleting")
-			return nil
+		if !forcePatch {
+			shouldSkip, forcePatch = askIfShouldSkipLibrary(finding.Path)
+			if !forcePatch && shouldSkip {
+				log.Info().
+					Str("findingPath", finding.Path).
+					Msg("Skipping library for patching")
+				continue
+			}
 		}
-		log.Debug().
-			Str("path", finding.Path).
-			Str("zipFilePath", finding.JndiLookupFileName).
-			Msg("Found file to remove")
+
+		err = patchJavaArchive(finding)
+		if err != nil {
+			continue
+		}
+		patchedLibraries = append(patchedLibraries, finding.Path)
 	}
 
+	log.Info().
+		Strs("patchedLibraries", patchedLibraries).
+		Msg("Successfully patched libraries.")
 	return nil
 }

tools/log4shell/commands/scan.go

@@ -76,6 +76,11 @@ func scanDirectoriesForVulnerableLibraries(
 	scanner := scan.NewLog4jDirectoryScanner(
 		excludeDirs, onlyScanArchives, noFollowSymlinks, processArchiveFile)
 
+	log.Info().
+		Strs("searchDirs", searchDirs).
+		Strs("excludeDirs", excludeDirs).
+		Msg("Scanning directories for vulnerable Log4j libraries.")
+
 	scannerFindings = scanner.Scan(searchDirs)
 	return
 }