Added hashselect arg option to hash() (Velocidex#1637)

* Added hashselect arg option to Hash vql function

* Added test and refactored code a bit

Co-authored-by: Mike Cohen <mike@velocidex.com>

artifacts/testdata/server/testcases/hash.in.yaml

@@ -0,0 +1,7 @@
+Queries:
+  - SELECT hash(path=srcDir+"/artifacts/testdata/files/hello.zip") AS AllHashes,
+     hash(path=srcDir+"/artifacts/testdata/files/hello.zip",
+          hashselect="md5") AS MD5HashOnly,
+     hash(path=srcDir+"/artifacts/testdata/files/hello.zip",
+          hashselect=["md5", "sha256"]) AS MDHashAndSha256
+    FROM scope()

artifacts/testdata/server/testcases/hash.out.yaml

@@ -0,0 +1,19 @@
+SELECT hash(path=srcDir+"/artifacts/testdata/files/hello.zip") AS AllHashes, hash(path=srcDir+"/artifacts/testdata/files/hello.zip", hashselect="md5") AS MD5HashOnly, hash(path=srcDir+"/artifacts/testdata/files/hello.zip", hashselect=["md5", "sha256"]) AS MDHashAndSha256 FROM scope()[
+ {
+  "AllHashes": {
+   "MD5": "cf2b61489d0564ec415caaa298c5cdd4",
+   "SHA1": "1f7e03e2b3b3225a2e3a9f54daef7d3799f2a4d6",
+   "SHA256": "3726e4d9af228eb0506913c15e0d2db20ec9c0b584d9d698a2254ef537dcf674"
+  },
+  "MD5HashOnly": {
+   "MD5": "cf2b61489d0564ec415caaa298c5cdd4",
+   "SHA1": "",
+   "SHA256": ""
+  },
+  "MDHashAndSha256": {
+   "MD5": "cf2b61489d0564ec415caaa298c5cdd4",
+   "SHA1": "",
+   "SHA256": "3726e4d9af228eb0506913c15e0d2db20ec9c0b584d9d698a2254ef537dcf674"
+  }
+ }
+]

docs/references/vql.yaml

@@ -1785,6 +1785,9 @@
   - name: accessor
     type: string
     description: The accessor to use
+  - name: hashselect
+    type: []string
+    description: The hash function to use (MD5,SHA1,SHA256)
   category: plugin
 - name: http_client
   description: |

vql/functions/hash.go

@@ -53,11 +53,12 @@ type HashResult struct {
 }
 
 type HashFunctionArgs struct {
-	Path     *accessors.OSPath `vfilter:"required,field=path,doc=Path to open and hash."`
-	Accessor string            `vfilter:"optional,field=accessor,doc=The accessor to use"`
+	Path       *accessors.OSPath `vfilter:"required,field=path,doc=Path to open and hash."`
+	Accessor   string            `vfilter:"optional,field=accessor,doc=The accessor to use"`
+	HashSelect []string          `vfilter:"optional,field=hashselect,doc=The hash function to use (MD5,SHA1,SHA256)"`
 }
 
-// The hash fuction calculates a hash of a file. It may be expensive
+// HashFunction calculates a hash of a file. It may be expensive
 // so we make it cancelllable.
 type HashFunction struct{}
 
@@ -95,10 +96,29 @@ func (self *HashFunction) Call(ctx context.Context,
 	}
 	defer file.Close()
 
-	result := HashResult{
-		md5:    md5.New(),
-		sha1:   sha1.New(),
-		sha256: sha256.New(),
+	result := HashResult{}
+
+	if arg.HashSelect == nil {
+		result = HashResult{
+			md5:    md5.New(),
+			sha1:   sha1.New(),
+			sha256: sha256.New(),
+		}
+	} else {
+		for _, hash_opt := range arg.HashSelect {
+			switch hash_opt {
+			case "sha256", "SHA256":
+				result.sha256 = sha256.New()
+			case "sha1", "SHA1":
+				result.sha1 = sha1.New()
+			case "md5", "MD5":
+				result.md5 = md5.New()
+			default:
+				scope.Log("hashselect option %s not recognized (should be md5, sha1, sha256)",
+					hash_opt)
+				return vfilter.Null{}
+			}
+		}
 	}
 
 	for {
@@ -112,24 +132,39 @@ func (self *HashFunction) Call(ctx context.Context,
 			// We are done!
 			if n == 0 || err == io.EOF {
 				if n == 0 {
-					result.MD5 = fmt.Sprintf(
-						"%x", result.md5.Sum(nil))
-					result.SHA1 = fmt.Sprintf(
-						"%x", result.sha1.Sum(nil))
-					result.SHA256 = fmt.Sprintf(
-						"%x", result.sha256.Sum(nil))
-
+					if result.md5 != nil {
+						result.MD5 = fmt.Sprintf(
+							"%x", result.md5.Sum(nil))
+					}
+
+					if result.sha1 != nil {
+						result.SHA1 = fmt.Sprintf(
+							"%x", result.sha1.Sum(nil))
+					}
+
+					if result.sha256 != nil {
+						result.SHA256 = fmt.Sprintf(
+							"%x", result.sha256.Sum(nil))
+					}
 					return result
 				}
 
 			} else if err != nil {
-				scope.Log(err.Error())
+				scope.Log("hash: %v", err)
 				return vfilter.Null{}
 			}
 
-			_, _ = result.md5.Write(buf[:n])
-			_, _ = result.sha1.Write(buf[:n])
-			_, _ = result.sha256.Write(buf[:n])
+			if result.md5 != nil {
+				_, _ = result.md5.Write(buf[:n])
+			}
+
+			if result.sha1 != nil {
+				_, _ = result.sha1.Write(buf[:n])
+			}
+
+			if result.sha256 != nil {
+				_, _ = result.sha256.Write(buf[:n])
+			}
 
 			// Charge an op for each buffer we read
 			scope.ChargeOp()
@@ -142,6 +177,7 @@ func (self HashFunction) Info(scope vfilter.Scope, type_map *vfilter.TypeMap) *v
 		Name:    "hash",
 		Doc:     "Calculate the hash of a file.",
 		ArgType: type_map.AddType(scope, &HashFunctionArgs{}),
+		Version: 2,
 	}
 }