Handle file not found correctly for Malware Bazaar.

Author: xorhex

download.go

@@ -57,6 +57,10 @@ type MalwareBazarQueryData struct {
 	File_name     string `json:"file_name"`
 }
 
+type MalwareBazaarQueryStatus struct {
+	Status string `json:"query_status"`
+}
+
 type AssemblyLineQuery struct {
 	Error_message  string                     `json:"api_error_message"`
 	Response       *AssemblyLineQueryResponse `json:"api_response"`
@@ -958,33 +962,43 @@ func malwareBazaarDownload(uri string, api string, hash Hash, doNotExtract bool,
 				fmt.Printf("    [!] Normally the response code: %s means that the provided URL %s needs a trailing slash (to avoid the redirect), but this already has a trailing slash.\nPlease file a bug report at https://github.com/xorhex/mlget/issues\n", response.Status, uri)
 			}
 		} else {
-			fmt.Printf("    [!] %s\n", response.Status)
-		}
-	}
+			byteValue, _ := io.ReadAll(response.Body)
 
-	err = writeToFile(response.Body, hash.Hash+".zip")
-	if err != nil {
-		fmt.Println(err)
-		return false, ""
-	}
+			var data = MalwareBazaarQueryStatus{}
+			error = json.Unmarshal(byteValue, &data)
 
-	fmt.Printf("    [+] Downloaded %s\n", hash.Hash+".zip")
-	if doNotExtract {
-		return true, hash.Hash + ".zip"
-	} else {
-		fmt.Println("    [-] Extracting...")
-		files, err := extractPwdZip(hash.Hash+".zip", password, true, hash)
-		if err != nil {
-			fmt.Println(err)
-			return false, ""
-		} else {
-			for _, f := range files {
-				fmt.Printf("    [-] Extracted %s\n", f.Name)
+			if error == nil {
+				if data.Status == "file_not_found" {
+					return false, ""
+				}
+			} else {
+				err = writeToFile(io.NopCloser(bytes.NewReader(byteValue)), hash.Hash+".zip")
+				if err != nil {
+					fmt.Println(err)
+					return false, ""
+				}
+
+				fmt.Printf("    [+] Downloaded %s\n", hash.Hash+".zip")
+				if doNotExtract {
+					return true, hash.Hash + ".zip"
+				} else {
+					fmt.Println("    [-] Extracting...")
+					files, err := extractPwdZip(hash.Hash+".zip", password, true, hash)
+					if err != nil {
+						fmt.Println(err)
+						return false, ""
+					} else {
+						for _, f := range files {
+							fmt.Printf("    [-] Extracted %s\n", f.Name)
+						}
+					}
+					os.Remove(hash.Hash + ".zip")
+					return true, hash.Hash
+				}
 			}
 		}
-		os.Remove(hash.Hash + ".zip")
-		return true, hash.Hash
 	}
+	return false, ""
 }
 
 func filescanio(uri string, api string, hash Hash, doNotExtract bool, password string) (bool, string) {

mlget-test-config/samples.yaml

@@ -63,4 +63,7 @@ test 22:
   hash: fc17c021f18ec73d1544ad46dde6a1f1949f126bf3e75f97e241f982e2b07c86
 test 23:
   name: TestVirusExchangeV2
-  hash: 0cacdb88b24bd34b9d8ef600b06814f76206e60e70f975c8e4bdaa1ab7cebb80 
+  hash: 0cacdb88b24bd34b9d8ef600b06814f76206e60e70f975c8e4bdaa1ab7cebb80 
+test 24:
+  name: TestMalwareBazaarNotFound
+  hash: fee889e9518d1c660bd6fa331c19aabada7eeff8f1c99f2ef4d64c662ed5805a 

mlget.go

@@ -34,7 +34,7 @@ var uploadToAssemblyLineFlag bool
 var uploadToAssemblyLineAndDeleteFlag bool
 var forceResubmission bool
 
-var version string = "3.4.2"
+var version string = "3.4.3"
 
 func usage() {
 	fmt.Println("mlget - A command line tool to download malware from a variety of sources")

mlget_test.go

@@ -509,6 +509,32 @@ func TestMalwareBazaar(t *testing.T) {
 	}
 }
 
+func TestMalwareBazaarNotFound(t *testing.T) {
+	home, _ := os.UserHomeDir()
+	cfg, err := LoadConfig(path.Join(home, ".mlget.yml"))
+	if err != nil {
+		log.Fatal()
+		t.Errorf("%v", err)
+	}
+
+	scfg, err := parseTestConfig("./mlget-test-config/samples.yaml", t.Name())
+	if err != nil {
+		log.Fatal()
+		t.Errorf("%v", err)
+	}
+
+	ht, _ := hashType(scfg.Hash)
+	hash := Hash{HashType: ht, Hash: scfg.Hash}
+
+	var osq ObjectiveSeeQuery
+	result, filename, _ := MalwareBazaar.QueryAndDownload(cfg, hash, false, osq)
+
+	if result {
+		os.Remove(filename)
+		t.Errorf("MalwareBazaar failed")
+	}
+}
+
 func TestMalpedia(t *testing.T) {
 	home, _ := os.UserHomeDir()
 	cfg, err := LoadConfig(path.Join(home, ".mlget.yml"))