fix for (s *stringObfuscator) Obfuscate()

[mh-cbon]

hey, i had to fix a small bug in the mentioned method.

The fix is

	if lastIndex < len(data) {
		result.Write(data[lastIndex:])
	}

The test program is trying to obfuscate itself and bugs with a panic: runtime error: slice bounds out of range

package main

var Appname = ""

func main() {
	log.Fatal(processPackage())
}

func processPackage() error {
	var opts packageProcessOpts

	flag.StringVar(&opts.Tag, "tag", Appname, "the output build tag of the generated files")

	flag.Parse()

	pkgpath := flag.Arg(0)

	if pkgpath == "" {
		return fmt.Errorf("invalid options, missing package. command line is %v <opt> <package>", Appname)
	}

	var conf loader.Config
	conf.Import(pkgpath)

	prog, err := conf.Load()
	if err != nil {
		return err
	}

	pkg := prog.Package(pkgpath)
	if pkg == nil {
		return fmt.Errorf("package %q not found in loaded program", pkgpath)
	}
	//-
	for _, f := range pkg.Files {
		b := new(bytes.Buffer)
		fset := token.NewFileSet()
		printer.Fprint(b, fset, f)

		obfuscator := &stringObfuscator{Contents: b.Bytes()}
		for _, decl := range f.Decls {
			ast.Walk(obfuscator, decl)
		}
		newCode, err := obfuscator.Obfuscate()
		if err != nil {
			return err
		}
		os.Stdout.Write(newCode)
		return nil
	}
	return nil
}

type stringObfuscator struct {
	Contents []byte
	Nodes    []*ast.BasicLit
}

func (s *stringObfuscator) Visit(n ast.Node) ast.Visitor {
	if lit, ok := n.(*ast.BasicLit); ok {
		if lit.Kind == token.STRING {
			s.Nodes = append(s.Nodes, lit)
		}
		return nil
	} else if decl, ok := n.(*ast.GenDecl); ok {
		if decl.Tok == token.CONST || decl.Tok == token.IMPORT {
			return nil
		}
	} else if _, ok := n.(*ast.StructType); ok {
		// Avoid messing with annotation strings.
		return nil
	}
	return s
}

func (s *stringObfuscator) Obfuscate() ([]byte, error) {
	sort.Sort(s)

	parsed := make([]string, s.Len())
	for i, n := range s.Nodes {
		var err error
		parsed[i], err = strconv.Unquote(n.Value)
		if err != nil {
			return nil, err
		}
	}

	var lastIndex int
	var result bytes.Buffer
	data := s.Contents
	for i, node := range s.Nodes {
		strVal := parsed[i]
		startIdx := node.Pos() - 1
		endIdx := node.End() - 1
		result.Write(data[lastIndex:startIdx])
		result.Write(obfuscatedStringCode(strVal))
		lastIndex = int(endIdx)
	}
	// if lastIndex < len(data) {
		result.Write(data[lastIndex:])
	// }
	return result.Bytes(), nil
}

func (s *stringObfuscator) Len() int {
	return len(s.Nodes)
}

func (s *stringObfuscator) Swap(i, j int) {
	s.Nodes[i], s.Nodes[j] = s.Nodes[j], s.Nodes[i]
}

func (s *stringObfuscator) Less(i, j int) bool {
	return s.Nodes[i].Pos() < s.Nodes[j].Pos()
}

func obfuscatedStringCode(str string) []byte {
	var res bytes.Buffer
	res.WriteString("(func() string {\n")
	res.WriteString("mask := []byte(\"")
	mask := make([]byte, len(str))
	for i := range mask {
		mask[i] = byte(rand.Intn(256))
		res.WriteString(fmt.Sprintf("\\x%02x", mask[i]))
	}
	res.WriteString("\")\nmaskedStr := []byte(\"")
	for i, x := range []byte(str) {
		res.WriteString(fmt.Sprintf("\\x%02x", x^mask[i]))
	}
	res.WriteString("\")\nres := make([]byte, ")
	res.WriteString(strconv.Itoa(len(mask)))
	res.WriteString(`)
        for i, m := range mask {
            res[i] = m ^ maskedStr[i]
        }
        return string(res)
        }())`)
	return res.Bytes()
}

[mh-cbon]

anyway after further inspection it appears it does not work properly at all, so maybe it is worth less to fix.

[mh-cbon]

I re open because after i start make use of astutil it just works!

see https://godoc.org/golang.org/x/tools/go/ast/astutil#Apply

func processPackage() error {
	var opts packageProcessOpts

	flag.StringVar(&opts.Tag, "tag", Appname, "the output build tag of the generated files")

	flag.Parse()

	pkgpath := flag.Arg(0)

	if pkgpath == "" {
		return fmt.Errorf("invalid options, missing package. command line is %v <opt> <package>", Appname)
	}

	var conf loader.Config
	conf.Import(pkgpath)

	prog, err := conf.Load()
	if err != nil {
		return err
	}

	pkg := prog.Package(pkgpath)
	if pkg == nil {
		return fmt.Errorf("package %q not found in loaded program", pkgpath)
	}
	//-
	for _, f := range pkg.Files {

		pre := func(c *astutil.Cursor) bool {
			if _, ok := c.Node().(*ast.ImportSpec); ok {
				return false
			}
			if x, ok := c.Node().(*ast.BasicLit); ok {
				b := obfuscatedStringCode(x.Value)
				expr, err := parser.ParseExpr(string(b))
				if err != nil {
					log.Println(err)
					return false
				}
				c.Replace(expr)
			}
			return true
		}
		res := astutil.Apply(f, pre, nil)

		log.Println("")
		log.Println(res)
		fset := token.NewFileSet()
		printer.Fprint(os.Stdout, fset, res)
	}
	return nil
}

[unixpickle]

Good find on astutil's Apply function! Regardless of all else, I should add a to-do to switch to that API for everything.

Do you have an example program that triggers the bug you found? I have not seen the error you mentioned before (the obfuscator has worked correctly on all programs I've tried).

I also noticed your new code likely doesn't work for string constants, e.g. const x = "hello world". Would be great to get it more fleshed out before merging.

[mh-cbon]

oh yeah it is possible i have forgotten some ast kinds. There are so many!

However, reading the earlier source code i think it should work for consts.
see this code https://play.golang.org/p/6jETLwwiYal the const become a basiclit within an ast.GenDecl which is not ast.ImportSpec so the program should jump into and visit the string node.

weird!

Anyways, sorry for earlier program example. I realized it was missing a type decl.

Here is an updated version which suits my need specifically.

I run the program against itself, something like echo "tomate" | go run main.go g -var Key github.com/... or go run main.go p -var Appname github.com/... where github... is the package path i m working in.

I checked for consts specifically, it worked here. But i have not written any test :x

type packageProcessOpts struct {
	Tag string
	Var string
}

func processPackage() error {
	var opts packageProcessOpts

	flag.StringVar(&opts.Tag, "tag", Appname, "the output build tag of the generated files")
	flag.StringVar(&opts.Var, "var", "", "only this var")

	flag.Parse()

	if opts.Tag == "" {
		return fmt.Errorf("-tag tagname is required")
	}

	pkgpath := flag.Arg(0)

	if pkgpath == "" {
		return fmt.Errorf("invalid options, missing package. command line is %v p <opt> <package>", Appname)
	}

	var conf loader.Config
	conf.Import(pkgpath)

	prog, err := conf.Load()
	if err != nil {
		return err
	}

	pkg := prog.Package(pkgpath)
	if pkg == nil {
		return fmt.Errorf("package %q not found in loaded program", pkgpath)
	}

	for _, f := range pkg.Files {
		pre := func(c *astutil.Cursor) bool {
			// log.Printf("%T\n", c.Node())
			if _, ok := c.Node().(*ast.ImportSpec); ok {
				return false
			}
			if x, ok := c.Node().(*ast.ValueSpec); ok {
				for i, n := range x.Names {
					if opts.Var == n.Name {
						v, ok := x.Values[i].(*ast.BasicLit)
						if ok {
							b := obfuscatedStringCode(v.Value)
							expr, err := parser.ParseExpr(string(b))
							if err != nil {
								log.Println(err)
							} else {
								x.Values[i] = expr
							}
						}
					}
				}
				return false
			}
			if x, ok := c.Node().(*ast.BasicLit); ok {
				if opts.Var == "" {
					b := obfuscatedStringCode(x.Value)
					expr, err := parser.ParseExpr(string(b))
					if err != nil {
						log.Println(err)
						return false
					}
					c.Replace(expr)
				}
			}
			return true
		}
		res := astutil.Apply(f, pre, nil)
		// os.Exit(0)
		fset := token.NewFileSet()
		fmt.Fprintf(os.Stdout, "//+build %v\n\n", opts.Tag)
		printer.Fprint(os.Stdout, fset, res)
	}
	return nil
}

type generateOpts struct {
	Tag     string
	Var     string
	Content string
}

func generate() error {
	var opts generateOpts

	flag.StringVar(&opts.Tag, "tag", Appname, "the output build tag of the generated files")
	flag.StringVar(&opts.Var, "var", "", "only this var")
	flag.StringVar(&opts.Content, "content", "", "the content of the var")

	flag.Parse()

	if opts.Tag == "" {
		return fmt.Errorf("-tag tagname is required")
	}
	if opts.Var == "" {
		return fmt.Errorf("-var tagname is required")
	}
	if opts.Content == "" {
		b := new(bytes.Buffer)
		done := make(chan error)
		go func() {
			_, err := io.Copy(b, os.Stdin)
			done <- err
		}()
		select {
		case err := <-done:
			if err != nil {
				return err
			}
		case <-time.After(time.Millisecond * 100):
			return fmt.Errorf("stdin is empty..")
		}
		opts.Content = b.String()
	}
	if opts.Content == "" {
		return fmt.Errorf("-content is required, use stdin otherwise")
	}

	pkgpath := flag.Arg(0)

	if pkgpath == "" {
		return fmt.Errorf("invalid options, missing package. command line is %v g <opt> <pkg path>", Appname)
	}

	var pkgName string

	var conf loader.Config
	conf.Import(pkgpath)

	prog, err := conf.Load()
	if err != nil {
		return err
	}

	pkg := prog.Package(pkgpath)
	if pkg == nil {
		return fmt.Errorf("package %q not found in loaded program", pkgpath)
	}

	if pkg.Pkg == nil {
		return fmt.Errorf("package %q not found in loaded program", pkgpath)
	}
	pkgName = pkg.Pkg.Name()
	if pkgName == "" {
		return fmt.Errorf("package name could not be determined")
	}

	obfuscatedStr := obfuscatedStringCode(opts.Content)

	filec := fmt.Sprintf(`package %v

var %v = %s
`, pkgName, opts.Var, obfuscatedStr)

	if opts.Tag != "" {
		filec = fmt.Sprintf(`//+build %v

%v`, opts.Tag, filec)
	}
	fmt.Println(filec)
	return nil
}

func obfuscatedStringCode(str string) []byte {
	var res bytes.Buffer
	res.WriteString("(func() string {\n")
	res.WriteString("mask := []byte(\"")
	mask := make([]byte, len(str))
	for i := range mask {
		mask[i] = byte(rand.Intn(256))
		res.WriteString(fmt.Sprintf("\\x%02x", mask[i]))
	}
	res.WriteString("\")\nmaskedStr := []byte(\"")
	for i, x := range []byte(str) {
		res.WriteString(fmt.Sprintf("\\x%02x", x^mask[i]))
	}
	res.WriteString("\")\nres := make([]byte, ")
	res.WriteString(strconv.Itoa(len(mask)))
	res.WriteString(`)
        for i, m := range mask {
            res[i] = m ^ maskedStr[i]
        }
        return string(res)
        }())`)
	return res.Bytes()
}