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()
}