perf: worker script

apps/behemoth/libs/go_script/Dockerfile

@@ -0,0 +1,16 @@
+FROM golang:1.21.12-alpine3.20
+
+WORKDIR /script
+
+COPY . .
+
+# ARG APT_MIRROR=mirrors.tuna.tsinghua.edu.cn
+ARG APT_MIRROR=mirrors.ustc.edu.cn
+# ARG APT_MIRROR=mirrors.aliyun.com
+RUN sed -i "[email protected]@${APT_MIRROR}@g" /etc/apk/repositories \
+    && apk update && apk add build-base
+
+ARG GOPROXY=https://goproxy.io
+RUN go mod tidy
+
+CMD ["/bin/sh"]

apps/behemoth/libs/go_script/build.sh

@@ -1,21 +1,22 @@
 #!/bin/bash
 
+set -e
+
 arch="amd64"
 output_name="jms_cli"
+IMAGE_NAME="jms_script"
+IMAGE_TAG="v1"
 #### Mac
-echo "1. 编译 Mac 版本脚本"
-bash -c "CGO_ENABLED=0 GOOS=darwin GOARCH=$arch go build -o ${output_name}_mac script.go"
-echo "编译 Mac 版本脚本成功"
-
-# 暂不支持
-#echo ""
-#### Windows
-#echo "2. 编译 Windows 版本脚本"
-#bash -c "CGO_ENABLED=0 GOOS=windows GOARCH=$arch go build -o ${output_name}_windows.exe script.go"
-#echo "编译 Windows 版本脚本成功"
+echo "编译 Mac 版本脚本"
+bash -c "CGO_ENABLED=1 GOOS=darwin GOARCH=$arch go build -o ${output_name}_darwin script.go"
 
 echo ""
 #### Linux
-echo "2. 编译 Linux 版本脚本"
-bash -c "CGO_ENABLED=0 GOOS=linux GOARCH=$arch go build -o ${output_name}_linux script.go"
-echo "编译 Linux 版本脚本成功"
+echo "编译 Linux 版本脚本"
+if docker images | grep -q "${IMAGE_NAME}[[:space:]]*${IMAGE_TAG}"; then
+  echo "镜像 ${IMAGE_NAME}:${IMAGE_TAG} 已经存在，跳过构建镜像"
+else
+  echo "镜像不存在，先构建镜像!"
+  bash -c "docker build -t jms_script:v1 ."
+fi
+bash -c "docker run --rm -v $(pwd):/script jms_script:v1 sh build_linux.sh"

apps/behemoth/libs/go_script/build_linux.sh

@@ -0,0 +1,9 @@
+#!/bin/sh
+
+set -e
+
+arch="amd64"
+output_name="jms_cli"
+#### Linux
+echo "编译 Linux 版本脚本"
+sh -c "CGO_ENABLED=1 GOOS=linux GOARCH=$arch go build -o ${output_name}_linux script.go"

apps/behemoth/libs/go_script/script.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/creack/pty"
 	_ "github.com/go-sql-driver/mysql"
+	_ "github.com/godror/godror"
 )
 
 const (
@@ -181,10 +182,49 @@ func (s *MySQLHandler) Close() {
 	_ = s.db.Close()
 }
 
+type OracleHandler struct {
+	opts CmdOptions
+
+	db *sql.DB
+}
+
+func (s *OracleHandler) Connect() error {
+	dsn := fmt.Sprintf(
+		`user="%s" password="%s" connectString="%s:%v/%s" sysdba=%v`,
+		s.opts.Auth.Username, s.opts.Auth.Password, s.opts.Auth.Address,
+		s.opts.Auth.Port, s.opts.Auth.DBName, s.opts.Auth.Privileged,
+	)
+	db, err := sql.Open("godror", dsn)
+	if err != nil {
+		return err
+	}
+	if err = db.Ping(); err != nil {
+		return err
+	}
+	s.db = db
+	return nil
+}
+
+func (s *OracleHandler) DoCommand(command string) (string, error) {
+	r, err := s.db.Exec(command)
+	if err != nil {
+		return "", err
+	}
+	affected, _ := r.RowsAffected()
+	return fmt.Sprintf("Affected rows: %v", affected), nil
+
+}
+
+func (s *OracleHandler) Close() {
+	_ = s.db.Close()
+}
+
 func getHandler(opts CmdOptions) BaseHandler {
 	switch opts.CmdType {
 	case "mysql":
 		return &MySQLHandler{opts: opts}
+	case "oracle":
+		return &OracleHandler{opts: opts}
 	case "script":
 		return &ScriptHandler{opts: opts}
 	}
@@ -199,15 +239,17 @@ type Cmd struct {
 }
 
 type Auth struct {
-	Address  string `json:"address"`
-	Port     int    `json:"port"`
-	Username string `json:"username"`
-	Password string `json:"password"`
-	DBName   string `json:"db_name"`
+	Address    string `json:"address"`
+	Port       int    `json:"port"`
+	Username   string `json:"username"`
+	Password   string `json:"password"`
+	DBName     string `json:"db_name"`
+	Privileged bool   `json:"privileged"`
 }
 
 type CmdOptions struct {
 	CommandBase64 string `json:"-"`
+	WithEnv       bool   `json:"-"`
 
 	TaskID      string   `json:"task_id"`
 	Host        string   `json:"host"`
@@ -470,6 +512,7 @@ func GetLogger(taskId string) *log.Logger {
 func main() {
 	opts := CmdOptions{}
 	flag.StringVar(&opts.CommandBase64, "command", opts.CommandBase64, "命令")
+	flag.BoolVar(&opts.WithEnv, "with_env", true, "使用参数中的环境变量执行脚本")
 	// 解析命令行标志
 	flag.Parse()
 
@@ -482,14 +525,18 @@ func main() {
 	jmsClient := NewJMSClient(opts.Host, opts.Token, opts.OrgId, logger)
 	logger.Printf("Start executing the task")
 
-	for _, part := range strings.Split(opts.Envs, ";") {
-		value := strings.Split(part, "=")
-		if len(value) == 2 {
-			if err := os.Setenv(value[0], value[1]); err != nil {
-				logger.Fatalf("Set environment variables failed: %v, error: %v", value, err)
-			} else {
-				logger.Printf("Set environment variables: %v=%v", value[0], value[1])
-			}
+	if opts.WithEnv {
+		envs := make([]string, 0)
+		for _, part := range strings.Split(opts.Envs, ";") {
+			envs = append(envs, part)
+			logger.Printf("Set environment variable %s", part)
+		}
+		cmd := exec.Command(os.Args[0], "--command", opts.CommandBase64)
+		cmd.Stdout = os.Stdout
+		cmd.Stderr = os.Stderr
+		cmd.Env = append(os.Environ(), envs...)
+		if err := cmd.Run(); err != nil {
+			log.Fatalf("failed to run command: %v", err)
 		}
 	}
 

apps/behemoth/libs/pools/worker.py

@@ -143,18 +143,23 @@ def __build_params(self, execution: Execution) -> dict:
         local_cmds_file = self.__generate_command_file(execution)
         remote_cmds_file = f'/tmp/behemoth/commands/{command_filepath}'
         token, __ = execution.user.create_bearer_token()
+        auth = {
+            'address': execution.asset.address,
+            'username': execution.account.username,
+            'password': execution.account.password,
+            'db_name': execution.asset.database.db_name
+        }
         if execution.asset.type == DatabaseTypes.MYSQL:
             cmd_type = script = 'mysql'
-            auth = {
-                'address': execution.asset.address,
+            auth['port'] = execution.asset.get_protocol_port('mysql'),
+        elif execution.asset.type == DatabaseTypes.ORACLE:
+            cmd_type = script = 'oracle'
+            auth.update({
                 'port': execution.asset.get_protocol_port('mysql'),
-                'username': execution.account.username,
-                'password': execution.account.password,
-                'db_name': execution.asset.database.db_name
-            }
+                'privileged': execution.account.privileged
+            })
         else:
             cmd_type = script = 'script'
-            auth = {}
         params: dict = {
             'host': settings.SITE_URL, 'cmd_type': cmd_type, 'script': script,
             'auth': auth, 'token': str(token), 'task_id': str(execution.id),

apps/behemoth/models.py

@@ -135,7 +135,7 @@ def _scp(self, local_path: str, remote_path: str, mode=0o544) -> None:
     def __ensure_script_exist(self) -> None:
         print(p.info('正在处理脚本文件'))
         platform_named = {
-            'mac': ('jms_cli_mac', '/tmp/behemoth', 'md5'),
+            'mac': ('jms_cli_darwin', '/tmp/behemoth', 'md5'),
             'linux': ('jms_cli_linux', '/tmp/behemoth', 'md5sum'),
             'windows': ('jms_cli_windows.exe', r'C:\Windows\Temp', ''),
         }