tiny-rdm/backend/utils/convert/pickle_convert.go

package convutil

import (
	"os/exec"
	"runtime"
)

type PickleConvert struct {
	CmdConvert
}

const pickleDecodeCode = `
import base64
import json
import pickle
import sys

if __name__ == "__main__":
    if len(sys.argv) >= 3:
        action = sys.argv[1].lower()
        content = sys.argv[2]

        try:
            if action == 'decode':
                decoded = base64.b64decode(content)
                obj = pickle.loads(decoded)
                unserialized = json.dumps(obj, ensure_ascii=False)
                print(base64.b64encode(unserialized.encode('utf-8')).decode('utf-8'))
            elif action == 'encode':
                decoded = base64.b64decode(content)
                obj = json.loads(decoded)
                serialized = pickle.dumps(obj)
                print(base64.b64encode(serialized).decode('utf-8'))
        except:
            print('[RDM-ERROR]')
    else:
        print('[RDM-ERROR]')

`

func NewPickleConvert() *PickleConvert {
	c := CmdConvert{
		Name: "Pickle",
		Auto: true,
	}
	c.DecodePath, c.EncodePath = "python3", "python3"
	var err error
	if _, err = exec.LookPath(c.DecodePath); err != nil {
		c.DecodePath, c.EncodePath = "python", "python"
		if _, err = exec.LookPath(c.DecodePath); err != nil {
			return nil
		}
	}
	// check if pickle available
	if runtime.GOOS == "darwin" {
		// the xcode-select installation prompt may appear on macOS
		// so check it manually in advance
		if _, err = exec.LookPath("xcode-select"); err != nil {
			return nil
		}
	}

	if _, err = runCommand(c.DecodePath, "-c", "import pickle"); err != nil {
		return nil
	}
	var filepath string
	if filepath, err = writeExecuteFile([]byte(pickleDecodeCode), "pickle_decoder.py"); err != nil {
		return nil
	}
	c.DecodeArgs = []string{filepath, "decode"}
	c.EncodeArgs = []string{filepath, "encode"}

	return &PickleConvert{
		CmdConvert: c,
	}
}

func (p *PickleConvert) Enable() bool {
	if p == nil {
		return false
	}
	return true
}

func (p *PickleConvert) Encode(str string) (string, bool) {
	if !p.Enable() {
		return str, false
	}
	return p.CmdConvert.Encode(str)
}

func (p *PickleConvert) Decode(str string) (string, bool) {
	if !p.Enable() {
		return str, false
	}
	return p.CmdConvert.Decode(str)
}