feat(threat): implement local datasets download & fallback to remote …

…repo (#168) * test(dsl): rm Log of Run results Signed-off-by: Dwi Siswanto <[email protected]> * feat(threat): implement local datasets download & fallback to remote repo Signed-off-by: Dwi Siswanto <[email protected]> * test(teler): add malformed tmp datasets cases Signed-off-by: Dwi Siswanto <[email protected]> * ci(tests): copy datasets to tmp on cache hit Signed-off-by: Dwi Siswanto <[email protected]> --------- Signed-off-by: Dwi Siswanto <[email protected]>
teler-sh · Feb 28, 2024 · 245d6fe · 245d6fe
1 parent 749ecc9
commit 245d6fe
Show file tree

Hide file tree

Showing 7 changed files with 165 additions and 17 deletions.
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
@@ -26,13 +26,16 @@ jobs:
         with:
           go-version: ${{ matrix.go-version }}
           cache-dependency-path: '**/go.sum'
-      - run: echo "date=$(date +%Y%m%d)" >> "$GITHUB_OUTPUT"
+      - run: echo "date=$(date +%d%m%Y)" >> "$GITHUB_OUTPUT"
         id: cache
       - uses: actions/cache@v4
+        id: teler-resources
         with:
           path: ~/.cache/teler-waf
           key: teler-resources-${{ steps.cache.outputs.date }}
           restore-keys: teler-resources-${{ steps.cache.outputs.date }}-
+      - run: cp -a ~/.cache/teler-waf/ /tmp/.teler-resources-${{ steps.cache.outputs.date }}
+        if: steps.teler-resources.outputs.cache-hit == 'true'
       - run: make ci
         if: (github.event_name != 'workflow_dispatch')
 

diff --git a/dsl/run_test.go b/dsl/run_test.go
@@ -30,9 +30,7 @@ func TestRunDSL(t *testing.T) {
 		program, err := env.Compile(`"1+10O`)
 		assert.NotNil(t, err)
 
-		res, err := env.Run(program)
+		_, err = env.Run(program)
 		assert.NotNil(t, err)
-
-		t.Log(res)
 	})
 }
diff --git a/go.mod b/go.mod
@@ -11,6 +11,7 @@ require (
 	github.com/go-playground/validator/v10 v10.16.0
 	github.com/hashicorp/go-getter v1.7.2
 	github.com/klauspost/compress v1.17.4
+	github.com/otiai10/copy v1.14.0
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/projectdiscovery/mapcidr v1.1.2
 	github.com/samber/lo v1.38.1
@@ -69,6 +70,7 @@ require (
 	go.uber.org/multierr v1.11.0 // indirect
 	golang.org/x/crypto v0.18.0 // indirect
 	golang.org/x/oauth2 v0.8.0 // indirect
+	golang.org/x/sync v0.3.0 // indirect
 	golang.org/x/sys v0.16.0 // indirect
 	golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect
 	google.golang.org/api v0.114.0 // indirect

diff --git a/go.sum b/go.sum
@@ -399,6 +399,9 @@ github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG
 github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
 github.com/mitchellh/go-testing-interface v1.14.1 h1:jrgshOhYAUVNMAJiKbEu7EqAwgJJ2JqpQmpLJOu07cU=
 github.com/mitchellh/go-testing-interface v1.14.1/go.mod h1:gfgS7OtZj6MA4U1UrDRp04twqAjfvlZyCfX3sDjEym8=
+github.com/otiai10/copy v1.14.0 h1:dCI/t1iTdYGtkvCuBG2BgR6KZa83PTclw4U5n2wAllU=
+github.com/otiai10/copy v1.14.0/go.mod h1:ECfuL02W+/FkTWZWgQqXPWZgW9oeKCSQ5qVfSc4qc4w=
+github.com/otiai10/mint v1.5.1 h1:XaPLeE+9vGbuyEHem1JNk3bYc7KKqyI/na0/mLd/Kks=
 github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
 github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
@@ -609,6 +612,8 @@ golang.org/x/sync v0.0.0-20220601150217-0de741cfad7f/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20220929204114-8fcdb60fdcc0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.3.0 h1:ftCYgMx6zT/asHUrPw8BLLscYtGznsLAnjq5RH9P66E=
+golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190312061237-fead79001313/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

diff --git a/teler_test.go b/teler_test.go
@@ -29,8 +29,7 @@ var (
 	})
 	client = &http.Client{}
 
-	homeDir  string
-	cacheDir = filepath.Join(".cache", "teler-waf")
+	cacheDir, tmpDir string
 )
 
 var mockRawReq = `POST /path?query=value#fragments HTTP/1.1
@@ -42,7 +41,17 @@ Content-Length: 9
 some=body`
 
 func init() {
-	homeDir, _ = os.UserHomeDir()
+	var err error
+
+	cacheDir, err = threat.Location()
+	if err != nil {
+		panic(err)
+	}
+
+	tmpDir, err = threat.TmpLocation()
+	if err != nil {
+		panic(err)
+	}
 
 	verified, err := threat.Verify()
 	if err != nil {
@@ -215,11 +224,12 @@ func TestNewWithWhitelist(t *testing.T) {
 }
 
 func TestNewWithMalformedDataset(t *testing.T) {
-	cvesPath := filepath.Join(homeDir, cacheDir, "cves.json")
+	cvesCachePath := filepath.Join(cacheDir, "cves.json")
+	cvesTmpPath := filepath.Join(tmpDir, "cves.json")
 
 	t.Run("nonexistent", func(t *testing.T) {
 		// Remove CVEs dataset
-		err := os.Remove(cvesPath)
+		err := os.Remove(cvesCachePath)
 		if err != nil && !os.IsNotExist(err) {
 			t.Fatal(err)
 		}
@@ -246,7 +256,7 @@ func TestNewWithMalformedDataset(t *testing.T) {
 
 	t.Run("malformed", func(t *testing.T) {
 		// Append CVEs dataset
-		f, err := os.OpenFile(cvesPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
+		f, err := os.OpenFile(cvesCachePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
 		if err != nil && !os.IsNotExist(err) {
 			t.Fatal(err)
 		}
@@ -275,6 +285,77 @@ func TestNewWithMalformedDataset(t *testing.T) {
 			t.Fatal(err)
 		}
 	})
+
+	t.Run("nonexistent-tmp", func(t *testing.T) {
+		// Remove CVEs dataset
+		err := os.Remove(cvesTmpPath)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// Remove cached datasets
+		err = os.RemoveAll(cacheDir)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// Initialize teler
+		telerMiddleware := New(Options{NoStderr: true})
+		wrappedHandler := telerMiddleware.Handler(handler)
+
+		// Create a test server with the wrapped handler
+		ts := httptest.NewServer(wrappedHandler)
+		defer ts.Close()
+
+		// Create a request to send to the test server
+		req, err := http.NewRequest("GET", ts.URL, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		_, err = client.Do(req)
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
+
+	t.Run("malformed-tmp", func(t *testing.T) {
+		// Append CVEs dataset
+		f, err := os.OpenFile(cvesTmpPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer f.Close()
+
+		if _, err := f.WriteString("AAAAAAAAAAAAAAAAAAaaaaaaaa\n"); err != nil {
+			t.Fatal(err)
+		}
+
+		// Remove cached datasets
+		err = os.RemoveAll(cacheDir)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// Initialize teler
+		telerMiddleware := New(Options{NoStderr: true})
+		wrappedHandler := telerMiddleware.Handler(handler)
+
+		// Create a test server with the wrapped handler
+		ts := httptest.NewServer(wrappedHandler)
+		defer ts.Close()
+
+		// Create a request to send to the test server
+		req, err := http.NewRequest("GET", ts.URL, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		_, err = client.Do(req)
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
 }
 
 func TestNewWithInMemory(t *testing.T) {

diff --git a/threat/const.go b/threat/const.go
@@ -5,6 +5,8 @@
 package threat
 
 const (
-	repoURL   = "https://github.com/kitabisa/teler-resources"
-	cachePath = "/teler-waf"
+	telerRes     = "teler-resources"
+	repoURL      = "https://github.com/kitabisa/" + telerRes
+	cachePath    = "/teler-waf"
+	tmpDirSuffix = "." + telerRes + "-%s"
 )
diff --git a/threat/get.go b/threat/get.go
@@ -12,6 +12,7 @@ import (
 	"path/filepath"
 
 	"github.com/hashicorp/go-getter"
+	"github.com/otiai10/copy"
 )
 
 // Get retrieves all the teler threat datasets.
@@ -38,15 +39,71 @@ func Get() error {
 		return err
 	}
 
-	// Retrieve the compressed archive DB file from the GitHub repository using go-getter
-	err = getter.Get(dst, fmt.Sprintf("%s?%s", DbURL, dbQuery))
+	// Downloading teler-resources from local, fallback to repo
+	if err := getLocal(); err != nil {
+		// Downloading from repo
+		err = getter.Get(dst, fmt.Sprintf("%s?%s", DbURL, dbQuery))
+		if err != nil {
+			return err
+		}
+
+		tmpDst, err := TmpLocation()
+		if err != nil {
+			return err
+		}
+
+		// Copy the resources from the cache dir to the temporary dir
+		err = copy.Copy(dst, tmpDst)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// getLocal fetches the local resources, copies them from a temporary location
+// to the cache directory, and verifies the checksum of the copied resources.
+// It returns an error if any step encounters an issue.
+func getLocal() error {
+	// Get the destination dir where local resources will be copied
+	dst, err := Location()
 	if err != nil {
-		// If there was an error retrieving the files, return the error
 		return err
 	}
 
-	// Return a nil error
-	return nil
+	// Get the temporary dir containing the resources to be copied
+	tmpDst, err := TmpLocation()
+	if err != nil {
+		return err
+	}
+
+	// Copy the resources from the temporary dir to the cache dir
+	err = copy.Copy(tmpDst, dst)
+	if err != nil {
+		return err
+	}
+
+	// Verify the checksum of the copied resources
+	_, err = Verify()
+
+	// Return any error encountered during the process
+	return err
+}
+
+// tmpLocation generates a temporary directory path based on the current date
+// and creates the directory if it doesn't already exist. It returns the
+// path of the temporary directory or an error if the creation fails.
+func TmpLocation() (string, error) {
+	date := time.Now().Format("02012006")
+	tmpDir := filepath.Join(os.TempDir(), fmt.Sprintf(tmpDirSuffix, date))
+
+	err := os.MkdirAll(tmpDir, 0755)
+	if err != nil && !os.IsExist(err) {
+		return "", err
+	}
+
+	return tmpDir, nil
 }
 
 // Location returns the location of the teler cache directory.