From 81078759d9e9fd5e3c140b1774538acf9256603f Mon Sep 17 00:00:00 2001 From: hari0205 Date: Thu, 20 Jul 2023 23:36:00 +0530 Subject: [PATCH 1/3] Test for api pkg --- go.mod | 4 ++ go.sum | 9 +++ internal/api_test.go | 162 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 175 insertions(+) create mode 100644 internal/api_test.go diff --git a/go.mod b/go.mod index b570991..66ce4a3 100644 --- a/go.mod +++ b/go.mod @@ -8,8 +8,12 @@ require ( ) require ( + github.com/davecgh/go-spew v1.1.1 // indirect github.com/edsrzf/mmap-go v1.1.0 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect github.com/saferwall/pe v1.4.4 // indirect + github.com/stretchr/testify v1.8.4 // indirect go.mozilla.org/pkcs7 v0.0.0-20210826202110-33d05740a352 // indirect golang.org/x/text v0.11.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index f31abb5..b1d3fc3 100644 --- a/go.sum +++ b/go.sum @@ -1,9 +1,15 @@ github.com/Binject/debug v0.0.0-20230508195519-26db73212a7a h1:4c0nc0krv8eh7gD809n+swLaCuFyHpxdrxwx0ZmHvBw= github.com/Binject/debug v0.0.0-20230508195519-26db73212a7a/go.mod h1:QzgxDLY/qdKlvnbnb65eqTedhvQPbaSP2NqIbcuKvsQ= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/edsrzf/mmap-go v1.1.0 h1:6EUwBLQ/Mcr1EYLE4Tn1VdW1A4ckqCQWZBw8Hr0kjpQ= github.com/edsrzf/mmap-go v1.1.0/go.mod h1:19H/e8pUPLicwkyNgOykDXkJ9F0MHE+Z52B8EIth78Q= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/saferwall/pe v1.4.4 h1:Ml++7/2/Z1iKwV4zCsd1nIqTEAdUQKAetwbbcCarhOg= github.com/saferwall/pe v1.4.4/go.mod h1:SNzv3cdgk8SBI0UwHfyTcdjawfdnN+nbydnEL7GZ25s= +github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= +github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= go.mozilla.org/pkcs7 v0.0.0-20210826202110-33d05740a352 h1:CCriYyAfq1Br1aIYettdHZTy8mBTIPo7We18TuO/bak= go.mozilla.org/pkcs7 v0.0.0-20210826202110-33d05740a352/go.mod h1:SNgMg+EgDFwmvSmLRTNKC5fegJjB7v23qTQ0XLGUNHk= @@ -35,3 +41,6 @@ golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGm golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal/api_test.go b/internal/api_test.go new file mode 100644 index 0000000..eb52113 --- /dev/null +++ b/internal/api_test.go @@ -0,0 +1,162 @@ +package pkg_test + +import ( + "fmt" + pkg "loldriverscan/internal" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" +) + +type testlolDrivers struct { + md5Map map[string]string + sha1Map map[string]string + sha256Map map[string]string + + authentihashMd5Map map[string]string + authentihashSha1Map map[string]string + authentihashSha256Map map[string]string + + lolDriverMap map[string]pkg.LolDriver +} + +type testDriverApi struct { + ID string + Category string + Samples []struct { + Filename string + MD5 string + Sha1 string + Sha256 string + OriginalFilename string + Authentihash pkg.Authentihash + } + Cve []string + CVEs []string +} + +func (tes *testlolDrivers) mockFindDriver(h pkg.Hashes, auth pkg.Authentihash) (pkg.LolDriver, error) { + id, ok := tes.authentihashMd5Map[auth.MD5] + if !ok { + id, ok = tes.authentihashSha1Map[auth.Sha1] + if !ok { + id, ok = tes.authentihashSha256Map[auth.Sha256] + if !ok { + // No matching driver found + return pkg.LolDriver{}, fmt.Errorf("no matching driver") + } + } + } + + // Found a matching driver ID, get the corresponding LolDriver object + lolDriver := tes.lolDriverMap[id] + + // Update the fields of the LolDriver object with the provided hashes + lolDriver.MD5 = h.Md5 + lolDriver.Sha1 = h.Sha1 + lolDriver.Sha256 = h.Sha256 + + // Update the fields of the Authentihash object + lolDriver.Authentihash.MD5 = auth.MD5 + lolDriver.Authentihash.Sha1 = auth.Sha1 + lolDriver.Authentihash.Sha256 = auth.Sha256 + + // Return the updated LolDriver object + return lolDriver, nil +} + +func mockCreateVulnerableDriverFinder() (*testlolDrivers, error) { + driverData, err := pkg.CreateVulnerableDriverFinder() +} + +func TestDrivers(t *testing.T) { + lolDriversApiUrl := `https://www.loldrivers.io/api/drivers.json` + + // Check if the the response status returned is 200 OK + res, err := http.Get(lolDriversApiUrl) + if err != nil { + t.Fatalf("Unable to make driver requestd: %v", err) + } + + if res.StatusCode != http.StatusOK { + t.Errorf("Expected status code to be %v but got %v", http.StatusOK, res.StatusCode) + } +} + +// func TestFetchApiNormaliseData(t *testing.T) { +// lolDriversApiUrl := `https://www.loldrivers.io/api/drivers.json` +// resp, err := http.Get(lolDriversApiUrl) +// if err != nil { +// t.Fatalf("[FATAL]Unable to get requested driver: %v", err) +// } +// defer resp.Body.Close() +// responseData, err := io.ReadAll(resp.Body) +// if err != nil { +// t.Fatalf("[FATAL]Unable to read data from response: %v", err) +// } + +// apiResponse := testDriverApi{} +// err = json.Unmarshal(responseData, &apiResponse) +// if err != nil { +// t.Error("Error unmarshalling response") +// } + +// } + +func TestFindDriver(t *testing.T) { + + //For test + testlolDrivers := testlolDrivers{ + authentihashMd5Map: map[string]string{ + "authMd5": "testdriver-1", + "authMD5hash": "testdriver-7", + }, + authentihashSha1Map: map[string]string{"authSha1": "testdriver-2"}, + authentihashSha256Map: map[string]string{"authSha256": "testdriver-3"}, + sha1Map: map[string]string{"hashSha1": "testdriver-4"}, + sha256Map: map[string]string{"hashSha256": "testdriver-5"}, + md5Map: map[string]string{"hashMd5": "testdriver-6"}, + lolDriverMap: map[string]pkg.LolDriver{ + "testdriver-1": {ID: "testdriver-1", Authentihash: pkg.Authentihash{MD5: "authMD5hash"}}, + "testdriver-2": {ID: "testdriver-2", Authentihash: pkg.Authentihash{Sha1: "authSha1hash"}}, + "testdriver-3": {ID: "testdriver-3", Authentihash: pkg.Authentihash{Sha256: "authSha256"}}, + "testdriver-4": {ID: "testdriver-4"}, + "testdriver-5": {ID: "testdriver-5"}, + "testdriver-6": {ID: "testdriver-6"}, + "testdriver-7": {ID: "testdriver-7", Authentihash: pkg.Authentihash{MD5: "authMd5"}}, + }, + } + + // Case 1: + t.Run("No matching drivers found", func(t *testing.T) { + _, err := testlolDrivers.mockFindDriver(pkg.Hashes{Md5: "testdriver-6"}, pkg.Authentihash{MD5: "authMD5hash"}) + if err == nil { + t.Errorf("Expected error, got nil error") + } + }) + + // Case 2: + t.Run("Valid drivers found", func(t *testing.T) { + foundDriver, err := testlolDrivers.mockFindDriver(pkg.Hashes{Md5: "testdriver-1"}, pkg.Authentihash{MD5: "authMd5"}) + if err != nil { + t.Errorf("Expected nil error, got error") + } + assert.Contains(t, foundDriver.Authentihash.MD5, "authMd5") + assert.Contains(t, foundDriver.MD5, "testdriver-1") + }) + + //Case 3: + t.Run("Correct drivers found", func(t *testing.T) { + foundDriver, err := testlolDrivers.mockFindDriver(pkg.Hashes{Md5: "testdriver-1"}, pkg.Authentihash{MD5: "authMD5hash"}) + fmt.Println(foundDriver) + if err != nil { + t.Errorf("Expected nil error, got error") + } + + }) +} + +func TestCreateVulnerableDriverFinder(t *testing.T) { + // +} From 9d2405f2a5c72987120c4dd4665be9c8e343f7f5 Mon Sep 17 00:00:00 2001 From: hari0205 Date: Sat, 22 Jul 2023 00:02:00 +0530 Subject: [PATCH 2/3] Added tests for internals --- go.mod | 2 + go.sum | 10 +++ internal/api_test.go | 97 ++++++++++---------- internal/drivers_test.go | 112 +++++++++++++++++++++++ internal/helpers_test.go | 186 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 356 insertions(+), 51 deletions(-) create mode 100644 internal/drivers_test.go create mode 100644 internal/helpers_test.go diff --git a/go.mod b/go.mod index 66ce4a3..ea63f2e 100644 --- a/go.mod +++ b/go.mod @@ -8,10 +8,12 @@ require ( ) require ( + bou.ke/monkey v1.0.2 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/edsrzf/mmap-go v1.1.0 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/saferwall/pe v1.4.4 // indirect + github.com/stretchr/objx v0.5.0 // indirect github.com/stretchr/testify v1.8.4 // indirect go.mozilla.org/pkcs7 v0.0.0-20210826202110-33d05740a352 // indirect golang.org/x/text v0.11.0 // indirect diff --git a/go.sum b/go.sum index b1d3fc3..2f5441b 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,8 @@ +bou.ke/monkey v1.0.2 h1:kWcnsrCNUatbxncxR/ThdYqbytgOIArtYWqcQLQzKLI= +bou.ke/monkey v1.0.2/go.mod h1:OqickVX3tNx6t33n1xvtTtu85YN5s6cKwVug+oHMaIA= github.com/Binject/debug v0.0.0-20230508195519-26db73212a7a h1:4c0nc0krv8eh7gD809n+swLaCuFyHpxdrxwx0ZmHvBw= github.com/Binject/debug v0.0.0-20230508195519-26db73212a7a/go.mod h1:QzgxDLY/qdKlvnbnb65eqTedhvQPbaSP2NqIbcuKvsQ= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/edsrzf/mmap-go v1.1.0 h1:6EUwBLQ/Mcr1EYLE4Tn1VdW1A4ckqCQWZBw8Hr0kjpQ= @@ -8,6 +11,12 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/saferwall/pe v1.4.4 h1:Ml++7/2/Z1iKwV4zCsd1nIqTEAdUQKAetwbbcCarhOg= github.com/saferwall/pe v1.4.4/go.mod h1:SNzv3cdgk8SBI0UwHfyTcdjawfdnN+nbydnEL7GZ25s= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= @@ -42,5 +51,6 @@ golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtn golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal/api_test.go b/internal/api_test.go index eb52113..0c41e45 100644 --- a/internal/api_test.go +++ b/internal/api_test.go @@ -21,19 +21,23 @@ type testlolDrivers struct { lolDriverMap map[string]pkg.LolDriver } -type testDriverApi struct { - ID string - Category string - Samples []struct { - Filename string - MD5 string - Sha1 string - Sha256 string - OriginalFilename string - Authentihash pkg.Authentihash +func mockfectchNormaliseData() ([]pkg.LolDriver, error) { + testdrivers := []pkg.LolDriver{ + { + ID: "driver-1", + Filename: "driver1.exe", + Path: "/sample/path/to/driver", + Status: "active", + Malicious: false, + MD5: "md5-hash", + Sha1: "sha1-hash", + Sha256: "sha256-hash", + CVEs: []string{"CVE-2021-1111", "CVE-2021-2222"}, + Authentihash: pkg.Authentihash{MD5: "auth-md5-1"}, + }, } - Cve []string - CVEs []string + return testdrivers, nil + } func (tes *testlolDrivers) mockFindDriver(h pkg.Hashes, auth pkg.Authentihash) (pkg.LolDriver, error) { @@ -66,10 +70,6 @@ func (tes *testlolDrivers) mockFindDriver(h pkg.Hashes, auth pkg.Authentihash) ( return lolDriver, nil } -func mockCreateVulnerableDriverFinder() (*testlolDrivers, error) { - driverData, err := pkg.CreateVulnerableDriverFinder() -} - func TestDrivers(t *testing.T) { lolDriversApiUrl := `https://www.loldrivers.io/api/drivers.json` @@ -84,25 +84,15 @@ func TestDrivers(t *testing.T) { } } -// func TestFetchApiNormaliseData(t *testing.T) { -// lolDriversApiUrl := `https://www.loldrivers.io/api/drivers.json` -// resp, err := http.Get(lolDriversApiUrl) -// if err != nil { -// t.Fatalf("[FATAL]Unable to get requested driver: %v", err) -// } -// defer resp.Body.Close() -// responseData, err := io.ReadAll(resp.Body) -// if err != nil { -// t.Fatalf("[FATAL]Unable to read data from response: %v", err) -// } - -// apiResponse := testDriverApi{} -// err = json.Unmarshal(responseData, &apiResponse) -// if err != nil { -// t.Error("Error unmarshalling response") -// } - -// } +func TestFetchApiNormaliseData(t *testing.T) { + sampledriver, err := mockfectchNormaliseData() // Mocking HTTP Get request + if err != nil { + t.Fail() // Will not occur + } + + assert.NotEmpty(t, sampledriver, "Driver list must not be empty") + assert.Equal(t, sampledriver[0].ID, "driver-1") +} func TestFindDriver(t *testing.T) { @@ -124,13 +114,20 @@ func TestFindDriver(t *testing.T) { "testdriver-4": {ID: "testdriver-4"}, "testdriver-5": {ID: "testdriver-5"}, "testdriver-6": {ID: "testdriver-6"}, - "testdriver-7": {ID: "testdriver-7", Authentihash: pkg.Authentihash{MD5: "authMd5"}}, }, } - // Case 1: - t.Run("No matching drivers found", func(t *testing.T) { - _, err := testlolDrivers.mockFindDriver(pkg.Hashes{Md5: "testdriver-6"}, pkg.Authentihash{MD5: "authMD5hash"}) + // Case 1.1: + t.Run("No matching drivers found [INVALID HASH]", func(t *testing.T) { + _, err := testlolDrivers.mockFindDriver(pkg.Hashes{Md5: "testdriver-1"}, pkg.Authentihash{MD5: "authMD5HASH"}) + if err == nil { + t.Errorf("Expected error, got nil error") + } + }) + + // Case 1.2: + t.Run("No matching drivers found [INVALID DRIVER]", func(t *testing.T) { + _, err := testlolDrivers.mockFindDriver(pkg.Hashes{Md5: "testdriver-2"}, pkg.Authentihash{Sha256: "hashSha256"}) if err == nil { t.Errorf("Expected error, got nil error") } @@ -142,21 +139,19 @@ func TestFindDriver(t *testing.T) { if err != nil { t.Errorf("Expected nil error, got error") } - assert.Contains(t, foundDriver.Authentihash.MD5, "authMd5") - assert.Contains(t, foundDriver.MD5, "testdriver-1") + assert.Equal(t, foundDriver.Authentihash.MD5, "authMd5") + assert.Equal(t, foundDriver.MD5, "testdriver-1") }) - //Case 3: - t.Run("Correct drivers found", func(t *testing.T) { - foundDriver, err := testlolDrivers.mockFindDriver(pkg.Hashes{Md5: "testdriver-1"}, pkg.Authentihash{MD5: "authMD5hash"}) - fmt.Println(foundDriver) - if err != nil { - t.Errorf("Expected nil error, got error") - } - - }) } func TestCreateVulnerableDriverFinder(t *testing.T) { - // + // Call the function + drivers, err := pkg.CreateVulnerableDriverFinder() + if err != nil { + t.Error("Expected nil error, got error") + } + + assert.NotEmpty(t, drivers, "Expected drivers to have values") + assert.NotNil(t, drivers, "Expected drivers to not be nil") } diff --git a/internal/drivers_test.go b/internal/drivers_test.go new file mode 100644 index 0000000..49d0366 --- /dev/null +++ b/internal/drivers_test.go @@ -0,0 +1,112 @@ +package pkg_test + +import ( + "errors" + "fmt" + pkg "loldriverscan/internal" + "testing" + + "bou.ke/monkey" + "github.com/stretchr/testify/assert" + "golang.org/x/sys/windows" + "golang.org/x/sys/windows/svc/mgr" +) + +func TestConnectToServiceManager(t *testing.T) { + mockWinHandle := windows.Handle(123) + + t.Run("ConnectToServiceManager without error", func(t *testing.T) { + patch := monkey.Patch(windows.OpenSCManager, func(host *uint16, database *uint16, access uint32) (windows.Handle, error) { + return mockWinHandle, nil + }) + + defer patch.Unpatch() + + mgr, err := pkg.ConnectToServiceManager() + if err != nil { + t.Error("Could not connect to service manager") + } + assert.NoError(t, err, "Unexpected error") + assert.NotNil(t, mgr, "Manager is nil") + assert.Equalf(t, mockWinHandle, mgr.Handle, "Unexpected handle value: Expected %v, got %v", mockWinHandle, mgr.Handle) + + }) + + t.Run("ConnectToServiceManager error", func(t *testing.T) { + mockwinError := windows.ERROR_ACCOUNT_RESTRICTION // Throws some error + patch := monkey.Patch(windows.OpenSCManager, func(host *uint16, database *uint16, access uint32) (windows.Handle, error) { + return mockWinHandle, mockwinError + }) + + defer patch.Unpatch() + + _, err := pkg.ConnectToServiceManager() + assert.Error(t, mockwinError, err) + assert.ErrorContainsf(t, err, "Account restrictions", "Expected error message %v, got %v", "Account restrictions", err) + + }) +} + +func TestListDriverServices(t *testing.T) { + + // Case 1: passing an invalid Mgr + t.Run("Invalid Mgr", func(t *testing.T) { + _, err := pkg.ListDriverServices(&mgr.Mgr{}) + assert.ErrorContains(t, err, "handle is invalid") + }) + + t.Run("Passing a valid service manager", func(t *testing.T) { + + patchListDriverService := monkey.Patch(pkg.ListDriverServices, func(m *mgr.Mgr) ([]string, error) { + if m.Handle == 123 { + return []string{"Sample service 1"}, nil + } else if m.Handle > 400 { + return []string{"Sample service 2"}, nil + } else { + return nil, fmt.Errorf("Invalid service manager handle") + } + + }) + defer patchListDriverService.Unpatch() + + res, err := pkg.ListDriverServices(&mgr.Mgr{Handle: 123}) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + assert.Containsf(t, res, "Sample service 1", "Expected service %v , got %v", "Sample service 1", res) + + }) + +} + +func TestOpenService(t *testing.T) { + OpenServicePatch := monkey.Patch(windows.OpenService, func(m windows.Handle, servicename *uint16, access uint32) (handle windows.Handle, err error) { + if m == 12345 { + return windows.Handle(12345), nil + } else { + return 0, fmt.Errorf("Invalid service Handle") + } + + }) + + defer OpenServicePatch.Unpatch() + + t.Run("Proper Service Openec", func(t *testing.T) { + mockOpenService, err := pkg.OpenService(&mgr.Mgr{Handle: 12345}, "Sample new service") + if err != nil { + t.Errorf("Unexpected error opening Handle %v", err) + } + + assert.Equalf(t, "Sample new service", mockOpenService.Name, "Expected service name %v but got %v:", "Sample new service", mockOpenService.Name) + assert.EqualValues(t, uintptr(12345), mockOpenService.Handle) + }) + + t.Run("Error new opening service", func(t *testing.T) { + mockOpenService, err := pkg.OpenService(&mgr.Mgr{Handle: 121}, "Sample new service 2") + if assert.Error(t, err) { + assert.Equal(t, errors.New("Invalid service Handle"), err) + } + assert.Nil(t, mockOpenService) + }) +} diff --git a/internal/helpers_test.go b/internal/helpers_test.go new file mode 100644 index 0000000..68c5b9b --- /dev/null +++ b/internal/helpers_test.go @@ -0,0 +1,186 @@ +package pkg_test + +import ( + "crypto/md5" + "crypto/sha1" + "crypto/sha256" + "fmt" + "io/fs" + "io/ioutil" + pkg "loldriverscan/internal" + "os" + "strings" + "testing" + + "bou.ke/monkey" + "github.com/stretchr/testify/assert" +) + +func ioCopy(dst *strings.Builder, src *os.File) { + buf := make([]byte, 1024) + for { + n, err := src.Read(buf) + if n > 0 { + dst.Write(buf[:n]) + } + if err != nil { + break + } + } +} +func TestHeuristicNormalisePath(t *testing.T) { + + // Table test approach + testCases := []struct { + input string + expectedPath string + expectedErr error + }{ + { + input: `system32\example.dll`, + expectedPath: `C:\Windows\System32\example.dll`, + expectedErr: nil, + }, + { + input: `systemRoot\file.txt`, + expectedPath: `C:\Windows\file.txt`, + expectedErr: nil, + }, + { + input: `??\some\path\file.txt`, + expectedPath: `some\path\file.txt`, + expectedErr: nil, + }, + // Error with Heuristic normalisation of path + // { + // input: `system32\NonExistentFile.txt`, + // expectedPath: "", + // expectedErr: fmt.Errorf("normalised path C:\\Windows\\system32\\NonExistentFile.txt does not exist: The system cannot find the file specified."), + // }, + } + + patch := monkey.Patch(os.Stat, func(name string) (fs.FileInfo, error) { + return nil, nil + }) + defer patch.Unpatch() + + for _, test := range testCases { + normalisedPath, err := pkg.HeuristicNormalisePath(test.input) + if (err != nil && test.expectedErr == nil) || (err == nil && test.expectedErr != nil) || (err != nil && test.expectedErr != nil && err.Error() != test.expectedErr.Error()) { + t.Errorf("Unexpected error: got %v, expected %v", err, test.expectedErr) + } + + // Check if the normalized path matches the expected path + if !strings.EqualFold(normalisedPath, test.expectedPath) { + t.Errorf("Mismatched path: got %v, expected %v", normalisedPath, test.expectedPath) + } + + } +} + +func TestHashFile(t *testing.T) { + t.Run("Test Hash for existing files", func(t *testing.T) { + + // Make a temporary file + tmpFile, err := ioutil.TempFile("", "test_file.txt") + if err != nil { + t.Fatalf("Error creating temporary file: %v", err) + } + defer os.Remove(tmpFile.Name()) + + // Write some content into it + if _, err = tmpFile.WriteString("This is a test temp file"); err != nil { + t.Fatalf("Error writing to temp file: %v", err) + } + + defer tmpFile.Close() + + hash, err := pkg.HashFile(tmpFile.Name()) + if err != nil { + t.Fatalf("Unexpected error while hashing file: %v", err) + } + + expectedMd5hash := fmt.Sprintf("%x", md5.Sum([]byte("This is a test temp file"))) + expectedSha1hash := fmt.Sprintf("%x", sha1.Sum([]byte("This is a test temp file"))) + expectedSha256hash := fmt.Sprintf("%x", sha256.Sum256(([]byte("This is a test temp file")))) + + if hash.Md5 != expectedMd5hash { + t.Errorf("MD5 hash mismatch. Expected: %s, Got: %s", expectedMd5hash, hash.Md5) + } + if hash.Sha1 != expectedSha1hash { + t.Errorf("SHA1 hash mismatch. Expected: %s, Got: %s", expectedSha1hash, hash.Sha1) + } + if hash.Sha256 != expectedSha256hash { + t.Errorf("SHA256 hash mismatch. Expected: %s, Got: %s", expectedSha256hash, hash.Sha256) + } + }) + + t.Run("Hashing non existing files", func(t *testing.T) { + _, err := pkg.HashFile("testfile_2.txt") + if err == nil { + t.Error("Expected an error for non-existing file, but got nil.") + } + }) +} + +/* + * // PE parser handling - Get Auth hash + */ + +func TestPrintLolDrivers(t *testing.T) { + + // Test drivers + drivers := []pkg.LolDriver{ + { + Filename: "Driver1.sys", + Path: "C:\\Windows\\System32\\drivers\\", + Status: "Loaded", + Malicious: false, + MD5: "fcd6aa0a8c3f9dfc8efb5f49298a1109", + ID: "Driver1", + CVEs: []string{"CVE-2022-1234", "CVE-2022-5678"}, + Authentihash: pkg.Authentihash{ + Sha256: "abcd1234abcd1234abcd1234abcd1234abcd1234abcd1234abcd1234abcd1234", + }, + }, + { + Filename: "MaliciousDriver.sys", + Path: "C:\\Malicious\\", + Status: "Loaded", + Malicious: true, + MD5: "jcnruadbaguvixhesoyam", + ID: "MaliciousDriver", + CVEs: []string{"CVE-2023-1111"}, + Authentihash: pkg.Authentihash{ + Sha256: "efgh5678efgh5678efgh5678efgh5678efgh5678efgh5678efgh5678efgh5678", + }, + }, + } + + // Capturing STDOUT + oldOutput := os.Stdout + r, w, _ := os.Pipe() + os.Stdout = w + + // Call the function being tested + pkg.PrintLolDrivers(drivers) + w.Close() + os.Stdout = oldOutput + + // Read the output from the buffer + var outputBuilder strings.Builder + ioCopy(&outputBuilder, r) + actualOutput := outputBuilder.String() + + fmt.Println(actualOutput) + + for _, driver := range drivers { + assert.Contains(t, actualOutput, driver.Filename) + assert.Contains(t, actualOutput, driver.Path) + assert.Contains(t, actualOutput, driver.MD5) + assert.Contains(t, actualOutput, driver.ID) + assert.Contains(t, actualOutput, "fcd6aa0a8c3f9dfc8efb5f49298a1109") + assert.Contains(t, actualOutput, "jcnruadbaguvixhesoyam") + + } +} From be1d4044ff06c73b14c71a0dcc48df62ed5821c4 Mon Sep 17 00:00:00 2001 From: hari0205 Date: Sat, 22 Jul 2023 00:09:39 +0530 Subject: [PATCH 3/3] Code cleanup and refactor --- internal/api_test.go | 3 +++ internal/drivers_test.go | 3 ++- internal/helpers_test.go | 2 -- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/internal/api_test.go b/internal/api_test.go index 0c41e45..dd8f3bd 100644 --- a/internal/api_test.go +++ b/internal/api_test.go @@ -9,6 +9,7 @@ import ( "github.com/stretchr/testify/assert" ) +// For testing purposes type testlolDrivers struct { md5Map map[string]string sha1Map map[string]string @@ -21,6 +22,7 @@ type testlolDrivers struct { lolDriverMap map[string]pkg.LolDriver } +// For testing purposes func mockfectchNormaliseData() ([]pkg.LolDriver, error) { testdrivers := []pkg.LolDriver{ { @@ -40,6 +42,7 @@ func mockfectchNormaliseData() ([]pkg.LolDriver, error) { } +// For testing purposes func (tes *testlolDrivers) mockFindDriver(h pkg.Hashes, auth pkg.Authentihash) (pkg.LolDriver, error) { id, ok := tes.authentihashMd5Map[auth.MD5] if !ok { diff --git a/internal/drivers_test.go b/internal/drivers_test.go index 49d0366..e9f9ff5 100644 --- a/internal/drivers_test.go +++ b/internal/drivers_test.go @@ -16,6 +16,7 @@ func TestConnectToServiceManager(t *testing.T) { mockWinHandle := windows.Handle(123) t.Run("ConnectToServiceManager without error", func(t *testing.T) { + // Monkey patching to mock implementation of windows.OpenSCmanager patch := monkey.Patch(windows.OpenSCManager, func(host *uint16, database *uint16, access uint32) (windows.Handle, error) { return mockWinHandle, nil }) @@ -33,7 +34,7 @@ func TestConnectToServiceManager(t *testing.T) { }) t.Run("ConnectToServiceManager error", func(t *testing.T) { - mockwinError := windows.ERROR_ACCOUNT_RESTRICTION // Throws some error + mockwinError := windows.ERROR_ACCOUNT_RESTRICTION // Make it throws some error patch := monkey.Patch(windows.OpenSCManager, func(host *uint16, database *uint16, access uint32) (windows.Handle, error) { return mockWinHandle, mockwinError }) diff --git a/internal/helpers_test.go b/internal/helpers_test.go index 68c5b9b..659afe0 100644 --- a/internal/helpers_test.go +++ b/internal/helpers_test.go @@ -172,8 +172,6 @@ func TestPrintLolDrivers(t *testing.T) { ioCopy(&outputBuilder, r) actualOutput := outputBuilder.String() - fmt.Println(actualOutput) - for _, driver := range drivers { assert.Contains(t, actualOutput, driver.Filename) assert.Contains(t, actualOutput, driver.Path)