feat: add tests

internal/core/runner/python/prescript.py

@@ -19,10 +19,29 @@ running_path = sys.argv[1]
 if not running_path:
     exit(-1)
 
+# get decrypt key
+key = sys.argv[2]
+if not key:
+    exit(-1)
+
+from base64 import b64decode
+key = b64decode(key)
+
 os.chdir(running_path)
 
 {{preload}}
 
 lib.DifySeccomp({{uid}}, {{gid}}, {{enable_network}})
 
-{{code}}
+code = b64decode("{{code}}")
+
+def decrypt(code, key):
+    key_len = len(key)
+    code_len = len(code)
+    code = bytearray(code)
+    for i in range(code_len):
+        code[i] = code[i] ^ key[i % key_len]
+    return bytes(code)
+
+code = decrypt(code, key)
+exec(code)

internal/core/runner/python/python.go

@@ -1,7 +1,9 @@
 package python
 
 import (
+	"crypto/rand"
 	_ "embed"
+	"encoding/base64"
 	"fmt"
 	"os"
 	"os/exec"
@@ -39,7 +41,7 @@ func (p *PythonRunner) Run(
 	configuration := static.GetDifySandboxGlobalConfigurations()
 
 	// initialize the environment
-	untrusted_code_path, err := p.InitializeEnvironment(code, preload, options)
+	untrusted_code_path, key, err := p.InitializeEnvironment(code, preload, options)
 	if err != nil {
 		return nil, nil, nil, err
 	}
@@ -53,7 +55,6 @@ func (p *PythonRunner) Run(
 		output_handler.SetAfterExitHook(func() {
 			os.RemoveAll(root_path)
 			os.Remove(root_path)
-			os.Remove(untrusted_code_path)
 		})
 
 		// create a new process
@@ -61,6 +62,7 @@ func (p *PythonRunner) Run(
 			configuration.PythonPath,
 			untrusted_code_path,
 			LIB_PATH,
+			key,
 		)
 		cmd.Env = []string{}
 
@@ -91,7 +93,7 @@ func (p *PythonRunner) Run(
 	return output_handler.GetStdout(), output_handler.GetStderr(), output_handler.GetDone(), nil
 }
 
-func (p *PythonRunner) InitializeEnvironment(code string, preload string, options *types.RunnerOptions) (string, error) {
+func (p *PythonRunner) InitializeEnvironment(code string, preload string, options *types.RunnerOptions) (string, string, error) {
 	if !checkLibAvaliable() {
 		// ensure environment is reversed
 		releaseLibBinary()
@@ -130,6 +132,25 @@ func (p *PythonRunner) InitializeEnvironment(code string, preload string, option
 		1,
 	)
 
+	// generate a random 512 bit key
+	key_len := 64
+	key := make([]byte, key_len)
+	_, err := rand.Read(key)
+	if err != nil {
+		return "", "", err
+	}
+
+	// encrypt the code
+	encrypted_code := make([]byte, len(code))
+	for i := 0; i < len(code); i++ {
+		encrypted_code[i] = code[i] ^ key[i%key_len]
+	}
+
+	// encode code using base64
+	code = base64.StdEncoding.EncodeToString(encrypted_code)
+	// encode key using base64
+	encoded_key := base64.StdEncoding.EncodeToString(key)
+
 	code = strings.Replace(
 		script,
 		"{{code}}",
@@ -138,14 +159,14 @@ func (p *PythonRunner) InitializeEnvironment(code string, preload string, option
 	)
 
 	untrusted_code_path := fmt.Sprintf("%s/tmp/%s.py", LIB_PATH, temp_code_name)
-	err := os.MkdirAll(path.Dir(untrusted_code_path), 0755)
+	err = os.MkdirAll(path.Dir(untrusted_code_path), 0755)
 	if err != nil {
-		return "", err
+		return "", "", err
 	}
 	err = os.WriteFile(untrusted_code_path, []byte(code), 0755)
 	if err != nil {
-		return "", err
+		return "", "", err
 	}
 
-	return untrusted_code_path, nil
+	return untrusted_code_path, encoded_key, nil
 }

tests/integration_tests/.gitignore

@@ -0,0 +1 @@
+logs

tests/integration_tests/conf/config.yaml

@@ -0,0 +1,13 @@
+app:
+  port: 8194
+  debug: True
+  key: dify-sandbox
+max_workers: 4
+max_requests: 50
+worker_timeout: 5
+python_path: /usr/local/bin/python3
+enable_network: True # please make sure there is no network risk in your environment
+proxy:
+  socks5: ''
+  http: ''
+  https: ''

tests/integration_tests/init.go

@@ -0,0 +1,17 @@
+package integrationtests_test
+
+import (
+	"github.com/langgenius/dify-sandbox/internal/core/runner/python"
+	"github.com/langgenius/dify-sandbox/internal/static"
+	"github.com/langgenius/dify-sandbox/internal/utils/log"
+)
+
+func init() {
+	static.InitConfig("conf/config.yaml")
+
+	// Test case for sys_fork
+	err := python.PreparePythonDependenciesEnv()
+	if err != nil {
+		log.Panic("failed to initialize python dependencies sandbox: %v", err)
+	}
+}

tests/integration_tests/malicious_test.go

@@ -0,0 +1,25 @@
+package integrationtests_test
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/langgenius/dify-sandbox/internal/core/runner/types"
+	"github.com/langgenius/dify-sandbox/internal/service"
+)
+
+func TestSysFork(t *testing.T) {
+	// Test case for sys_fork
+	resp := service.RunPython3Code(`
+import os
+os.fork()
+	`, "", &types.RunnerOptions{})
+
+	if resp.Code != 0 {
+		t.Error(resp)
+	}
+
+	if !strings.Contains(resp.Data.(*service.RunCodeResponse).Stderr, "operation not permitted") {
+		t.Error(resp.Data.(*service.RunCodeResponse).Stderr)
+	}
+}

tests/integration_tests/ordinary_feature_test.go

@@ -0,0 +1,11 @@
+package integrationtests_test
+
+import (
+	"testing"
+
+	"github.com/langgenius/dify-sandbox/internal/static"
+)
+
+func TestBase64(t *testing.T) {
+	static.InitConfig("conf/config.yaml")
+}