From 23e9ac3d6fb474e82c87a244a28f5ea9a52a2b7e Mon Sep 17 00:00:00 2001
From: Mark Woan <markwoan@gmail.com>
Date: Tue, 13 Sep 2016 08:47:45 +0100
Subject: [PATCH] Fixed some minor bugs :-)

---
 source/src/woanware/lookuper/main.go    |  2 +-
 source/src/woanware/lookuper/vt_hash.go | 19 ++++++--
 source/src/woanware/lookuper/vt_url.go  | 65 ++++++++++++++++---------
 3 files changed, 58 insertions(+), 28 deletions(-)

diff --git a/source/src/woanware/lookuper/main.go b/source/src/woanware/lookuper/main.go
index 984b89c..9682555 100644
--- a/source/src/woanware/lookuper/main.go
+++ b/source/src/woanware/lookuper/main.go
@@ -29,7 +29,7 @@ var (
 // ##### Constants ############################################################
 
 const APP_NAME string = "lookuper"
-const APP_VERSION string = "0.0.1"
+const APP_VERSION string = "0.0.2"
 const DB_FILE_NAME string = "./lookuper.db"
 const CONFIG_FILE_NAME string = "./lookuper.config"
 
diff --git a/source/src/woanware/lookuper/vt_hash.go b/source/src/woanware/lookuper/vt_hash.go
index b1de46e..391672d 100644
--- a/source/src/woanware/lookuper/vt_hash.go
+++ b/source/src/woanware/lookuper/vt_hash.go
@@ -30,7 +30,14 @@ type VtHash struct {
 // Processes a VT API request for multiple hashes
 func(h *VtHash) Process(data []string) int8 {
 
-	frr, err := h.govtc.GetFileReports(data)
+	var err error
+	var fr *govt.FileReport
+	var frr *govt.FileReportResults
+	if len(data) == 1 {
+		fr, err = h.govtc.GetFileReport(data[0])
+	} else {
+		frr, err = h.govtc.GetFileReports(data)
+	}
 
 	if err != nil {
 		if strings.Contains(strings.ToLower(err.Error()), "unexpected status code: 204") {
@@ -41,9 +48,13 @@ func(h *VtHash) Process(data []string) int8 {
 		return WORK_RESPONSE_ERROR
 	}
 
-	for _, fr := range *frr {
-		if fr.ResponseCode == 1 {
-			h.setRecord(fr)
+	if len(data) == 1 {
+		h.setRecord(*fr)
+	} else {
+		for _, fr := range *frr {
+			if fr.ResponseCode == 1 {
+				h.setRecord(fr)
+			}
 		}
 	}
 
diff --git a/source/src/woanware/lookuper/vt_url.go b/source/src/woanware/lookuper/vt_url.go
index 0a19f93..29d126d 100644
--- a/source/src/woanware/lookuper/vt_url.go
+++ b/source/src/woanware/lookuper/vt_url.go
@@ -30,39 +30,55 @@ type VtUrl struct {
 
 // Processes a VT API request for a URL(s)
 func (u *VtUrl) Process(data []string) int8 {
-	//if isSingleItem == true {
-	//	ur, err := u.govtc.GetUrlReport(data[0])
-	//	if err != nil {
-	//		if strings.Contains(strings.ToLower(err.Error()), "unexpected status code: 204") {
-	//			return WORK_RESPONSE_KEY_FAILED
-	//		}
-	//
-	//		log.Printf("Error requesting VT URL report: %v", err)
-	//		return WORK_RESPONSE_ERROR
-	//	}
-	//
-	//	if ur.ResponseCode == 1 {
-	//		u.processUrlReport(ur)
-	//	}
-	//} else {
-		urr, err := u.govtc.GetUrlReports(data)
-		if err != nil {
-			if strings.Contains(strings.ToLower(err.Error()), "unexpected status code: 204") {
-				return WORK_RESPONSE_KEY_FAILED
-			}
 
-			log.Printf("Error requesting VT MD5 report: %v", err)
-			return WORK_RESPONSE_ERROR
+	var err error
+	var ur *govt.UrlReport
+	var urr *govt.UrlReports
+	if len(data) == 1 {
+		ur, err = u.govtc.GetUrlReport(data[0])
+	} else {
+		urr, err = u.govtc.GetUrlReports(data)
+	}
+
+	if err != nil {
+		if strings.Contains(strings.ToLower(err.Error()), "unexpected status code: 204") {
+			return WORK_RESPONSE_KEY_FAILED
 		}
 
+		log.Printf("Error requesting VT URL report: %v", err)
+		return WORK_RESPONSE_ERROR
+	}
+
+	if len(data) == 1 {
+		u.setRecord(*ur)
+	} else {
 		for _, ur := range *urr {
 			if ur.ResponseCode == 1 {
 				u.setRecord(ur)
 			}
 		}
-	//}
+	}
 
 	return WORK_RESPONSE_OK
+
+	//	urr, err := u.govtc.GetUrlReports(data)
+	//	if err != nil {
+	//		if strings.Contains(strings.ToLower(err.Error()), "unexpected status code: 204") {
+	//			return WORK_RESPONSE_KEY_FAILED
+	//		}
+	//
+	//		log.Printf("Error requesting VT MD5 report: %v", err)
+	//		return WORK_RESPONSE_ERROR
+	//	}
+	//
+	//	for _, ur := range *urr {
+	//		if ur.ResponseCode == 1 {
+	//			u.setRecord(ur)
+	//		}
+	//	}
+	////}
+	//
+	//return WORK_RESPONSE_OK
 }
 
 //
@@ -78,6 +94,7 @@ func (u *VtUrl) DoesDataExist(data string, staleTimestamp time.Time) (error, boo
 
 // Inserts a new URL record, if that fails due to it already existing, then retrieve details and update
 func (u *VtUrl) setRecord(ur govt.UrlReport) int8 {
+
 	data := new(VtUrl)
 	u.updateObject(data, ur)
 
@@ -107,6 +124,7 @@ func (u *VtUrl) setRecord(ur govt.UrlReport) int8 {
 
 // Generic method to copy the VT data to our URL object
 func (u *VtUrl) updateObject(url *VtUrl, ur govt.UrlReport) {
+
 	url.Url = ur.Resource
 	url.UrlMd5 = strings.ToLower(util.Md5HashString(ur.Resource))
 	url.Positives = int16(ur.Positives)
@@ -126,6 +144,7 @@ func (u *VtUrl) updateObject(url *VtUrl, ur govt.UrlReport) {
 
 // Creates a comma delimited string with the scan engine and the result/malware/virus
 func (u *VtUrl) generateUrlScansString(fs map[string]govt.UrlScan) string {
+
 	// We need to sort the keys first, since the iteration is actually random if not
 	var keys []string
 	for e, s := range fs {