Can wait even if file/url scan is in progress (#8)

* added waiting for files/urls that are already in scan

* hash search refactoring

README.md

@@ -59,7 +59,7 @@ At the moment `virusgotal` supports files/URLs scan and search results by file h
   `virusgotal url <URL> --wait`
   * Get JSON formatted result:
   `virusgotal url <FILE> --json`
-  
+
 * To lookup a hash sum in VirusTotal database:
 `virusgotal hash <HASH>`
   * Get JSON formatted result:
@@ -72,4 +72,4 @@ You also can combine options: `virusgotal file --wait --json --force <FILE>`
 All kinds of contribution are welcome! Please send your PRs, bug reports, ideas, suggestions.
 
 ## License
-MIT License
+MIT License

filescan.go

@@ -55,6 +55,24 @@ func printFileResult(result *govt.FileReport) {
 	}
 }
 
+// Invoke spinner to display scan progress
+func waitFileResult(vt *govt.Client, filename string) {
+	for m := 0; m <= 600; m += 30 { // m == minutes
+		loader(fmt.Sprintf("waiting for results for %d seconds", m))
+		r, err := vt.GetFileReport(sha256sum(filename))
+		check(err)
+		if r.Status.ResponseCode == 1 {
+			if !*jsonFile && !*jsonHash {
+				fmt.Printf("scan took ~ %d seconds\n", m)
+			}
+			printFileResult(r)
+		}
+	}
+	color.Set(color.FgHiRed)
+	fmt.Printf("\nSomething went wrong with VirusTotal. Maybe their servers are overloaded.\n" +
+		"Please try again later\n")
+}
+
 func scanFile(filename string) {
 	// Init VT
 	apikey := os.Getenv("VT_API_KEY")
@@ -71,10 +89,14 @@ func scanFile(filename string) {
 		case 1: // Results exist
 			printFileResult(r)
 		case -2: // Scan in progress
-			color.Set(color.FgHiRed)
-			fmt.Printf("You scan is still in progress\n")
-			color.Unset()
-			os.Exit(1)
+			if !*waitFile {
+				color.Set(color.FgHiRed)
+				fmt.Printf("You scan is still in progress\n")
+				color.Unset()
+				os.Exit(1)
+			} else {
+				waitFileResult(vt, filename)
+			}
 		}
 	}
 
@@ -90,19 +112,6 @@ func scanFile(filename string) {
 		color.Unset()
 	}
 	if *waitFile { // Wait for results if user wishes
-		for m := 0; m <= 600; m += 30 { // m == minutes
-			loader(fmt.Sprintf("waiting for results for %d seconds", m))
-			r, err := vt.GetFileReport(sha256sum(filename))
-			check(err)
-			if r.Status.ResponseCode == 1 {
-				if !*jsonFile && !*jsonHash {
-					fmt.Printf("scan took ~ %d seconds\n", m)
-				}
-				printFileResult(r)
-			}
-		}
-		color.Set(color.FgHiRed)
-		fmt.Printf("\nSomething went wrong with VirusTotal. Maybe their servers are overloaded.\n" +
-			"Please try again later\n")
+		waitFileResult(vt, filename)
 	}
 }

hashsearch.go

@@ -17,18 +17,19 @@ func searchHash(hash string) {
 	r, err := vt.GetFileReport(hash)
 	check(err)
 
-	if r.ResponseCode == 0 { // Hash not found
+	switch r.ResponseCode {
+	case 0: // Hash not found
 		color.Set(color.FgHiRed, color.Bold)
-		fmt.Printf("Given hash isn't recognized by VirusTotal\n")
+		fmt.Printf("\nGiven hash is not recognized by VirusTotal\n")
 		color.Unset()
 		os.Exit(1)
-	}
-	if r.ResponseCode == -2 { // File scan with given hash is still in progress
+	case -2: // File scan with given hash is still in progress
 		color.Set(color.FgHiYellow)
 		fmt.Printf("\nScan with given hash is still in progress\n")
 		color.Unset()
 		os.Exit(1)
 	}
+
 	if r.Positives > 0 { // Malware detected
 		if !*jsonHash {
 			color.Set(color.FgHiRed, color.Bold)

urlscan.go

@@ -36,6 +36,23 @@ func printUrlResult(result *govt.UrlReport) {
 	}
 }
 
+func waitUrlResult(vt *govt.Client, urlname string) {
+	for m := 0; m <= 600; m += 30 {
+		loader(fmt.Sprintf("waiting for results for %d seconds", m))
+		r, err := vt.GetUrlReport(urlname)
+		check(err)
+		if r.Status.ResponseCode == 1 {
+			if !*jsonUrl {
+				fmt.Printf("scan took ~ %d seconds\n", m)
+			}
+			printUrlResult(r)
+		}
+	}
+	color.Set(color.FgHiRed)
+	fmt.Printf("\nSomething went wrong with VirusTotal. Maybe their servers are overloaded.\n" +
+		"Please try again later\n")
+}
+
 func scanUrl(urlname string) {
 	// Init VT
 	apikey := os.Getenv("VT_API_KEY")
@@ -52,15 +69,18 @@ func scanUrl(urlname string) {
 		case 1: // Results exist
 			printUrlResult(r)
 		case -2: // Scan in progress
-			color.Set(color.FgHiRed)
-			fmt.Printf("You scan is still in progress\n")
-			color.Unset()
-			os.Exit(1)
+			if !*waitUrl {
+				color.Set(color.FgHiRed)
+				fmt.Printf("You scan is still in progress\n")
+				color.Unset()
+				os.Exit(1)
+			} else {
+				waitUrlResult(vt, urlname)
+			}
 		}
 	}
 
 	// Scan URL
-	//if !*waitUrl {
 	report, err := vt.ScanUrl(urlname)
 	check(err)
 	if !*jsonUrl {
@@ -72,19 +92,6 @@ func scanUrl(urlname string) {
 		color.Unset()
 	}
 	if *waitUrl { // Wait for results if user wishes
-		for m := 0; m <= 600; m += 30 {
-			loader(fmt.Sprintf("waiting for results for %d seconds", m))
-			r, err := vt.GetUrlReport(urlname)
-			check(err)
-			if r.Status.ResponseCode == 1 {
-				if !*jsonUrl {
-					fmt.Printf("scan took ~ %d seconds\n", m)
-				}
-				printUrlResult(r)
-			}
-		}
-		color.Set(color.FgHiRed)
-		fmt.Printf("\nSomething went wrong with VirusTotal. Maybe their servers are overloaded.\n" +
-			"Please try again later\n")
+		waitUrlResult(vt, urlname)
 	}
 }