From 971e595d3470661086cc88805b7ffe2209c7f17b Mon Sep 17 00:00:00 2001
From: Pagran <67878280+pagran@users.noreply.github.com>
Date: Wed, 12 Aug 2020 19:15:49 +0300
Subject: [PATCH] Add line obfuscator Add tiny flag

---
 line_obfuscator.go         | 42 ++++++++++++++++++++++++++++++++++++++
 main.go                    | 20 ++++++++++++++++--
 testdata/scripts/basic.txt |  2 +-
 3 files changed, 61 insertions(+), 3 deletions(-)
 create mode 100644 line_obfuscator.go

diff --git a/line_obfuscator.go b/line_obfuscator.go
new file mode 100644
index 0000000..e5db9f8
--- /dev/null
+++ b/line_obfuscator.go
@@ -0,0 +1,42 @@
+package main
+
+import (
+	"fmt"
+	"go/ast"
+	"golang.org/x/tools/go/ast/astutil"
+	mathrand "math/rand"
+)
+
+// PosMax is the largest line or column value that can be represented without loss.
+// Incoming values (arguments) larger than PosMax will be set to PosMax.
+// Source: https://golang.org/src/cmd/compile/internal/syntax/pos.go
+const PosMax = 1 << 30
+const PosMin = 1
+
+var emptyLine = &ast.CommentGroup{List: []*ast.Comment{{Text: "//line :1"}}}
+
+func transformLineInfo(fileIndex int, file *ast.File) *ast.File {
+	pre := func(cursor *astutil.Cursor) bool {
+		funcDecl, ok := cursor.Node().(*ast.FuncDecl)
+		if !ok {
+			return true
+		}
+
+		// Ignore functions with //go: directives
+		if funcDecl.Doc != nil && len(funcDecl.Doc.List) != 0 {
+			return true
+		}
+
+		if envGarbleTiny {
+			funcDecl.Doc = emptyLine
+			return true
+		}
+
+		line := hashWithAsUint64(buildInfo.buildID, fmt.Sprintf("%d:%s", fileIndex, funcDecl.Name), PosMin, PosMax)
+		comment := &ast.Comment{Text: fmt.Sprintf("//line %c.go:%d", nameCharset[mathrand.Intn(len(nameCharset))], line)}
+		funcDecl.Doc = &ast.CommentGroup{List: []*ast.Comment{comment}}
+		return true
+	}
+
+	return astutil.Apply(file, pre, nil).(*ast.File)
+}
diff --git a/main.go b/main.go
index c3d2cb7..79d4cad 100644
--- a/main.go
+++ b/main.go
@@ -37,6 +37,7 @@ var flagSet = flag.NewFlagSet("garble", flag.ContinueOnError)
 
 var (
 	flagGarbleLiterals bool
+	flagGarbleTiny     bool
 	flagDebugDir       string
 	flagSeed           string
 )
@@ -44,6 +45,7 @@ var (
 func init() {
 	flagSet.Usage = usage
 	flagSet.BoolVar(&flagGarbleLiterals, "literals", false, "Encrypt all literals with AES, currently only literal strings are supported")
+	flagSet.BoolVar(&flagGarbleTiny, "tiny", false, "Encrypt all literals with AES, currently only literal strings are supported")
 	flagSet.StringVar(&flagDebugDir, "debugdir", "", "Write the garbled source to a given directory: '-debugdir=./debug'")
 	flagSet.StringVar(&flagSeed, "seed", "", "Provide a custom base64-encoded seed: '-seed=o9WDTZ4CN4w=' \nFor a random seed provide: '-seed=random'")
 }
@@ -72,7 +74,8 @@ var (
 	deferred []func() error
 	fset     = token.NewFileSet()
 
-	b64         = base64.NewEncoding("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789_z")
+	nameCharset = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789_z"
+	b64         = base64.NewEncoding(nameCharset)
 	printConfig = printer.Config{Mode: printer.RawFormat}
 
 	// listPackage helps implement a types.Importer which finds the export
@@ -94,6 +97,7 @@ var (
 
 	envGarbleDir      = os.Getenv("GARBLE_DIR")
 	envGarbleLiterals = os.Getenv("GARBLE_LITERALS") == "true"
+	envGarbleTiny     = os.Getenv("GARBLE_TINY") == "true"
 	envGarbleDebugDir = os.Getenv("GARBLE_DEBUGDIR")
 	envGarbleSeed     = os.Getenv("GARBLE_SEED")
 	envGoPrivate      string // filled via 'go env' below to support 'go env -w'
@@ -247,6 +251,7 @@ func mainErr(args []string) error {
 		}
 		os.Setenv("GARBLE_DIR", wd)
 		os.Setenv("GARBLE_LITERALS", fmt.Sprint(flagGarbleLiterals))
+		os.Setenv("GARBLE_TINY", fmt.Sprint(flagGarbleTiny))
 
 		if flagSeed == "random" {
 			seed = make([]byte, 16) // random 128 bit seed
@@ -457,7 +462,7 @@ func transformCompile(args []string) ([]string, error) {
 		return nil, err
 	}
 	deferred = append(deferred, func() error {
-		return os.RemoveAll(tempDir)
+		return nil //return os.RemoveAll(tempDir)
 	})
 
 	// Add our temporary dir to the beginning of -trimpath, so that we don't
@@ -504,6 +509,7 @@ func transformCompile(args []string) ([]string, error) {
 			// messy.
 			name = "_cgo_" + name
 		default:
+			file = transformLineInfo(i, file)
 			file = transformGo(file, info, blacklist)
 			name = fmt.Sprintf("z%d.go", i)
 
@@ -657,6 +663,16 @@ func hashWith(salt, value string) string {
 	return "z" + sum[:length]
 }
 
+func hashWithAsUint64(salt, value string, min, max uint64) uint64 {
+	d := sha256.New()
+	io.WriteString(d, salt)
+	d.Write(seed)
+	io.WriteString(d, value)
+	sum := d.Sum(nil)
+	val := binary.LittleEndian.Uint64(sum)
+	return min + (val % (max - min))
+}
+
 // buildBlacklist collects all the objects in a package which are known to be
 // used with reflect.TypeOf or reflect.ValueOf. Since we obfuscate one package
 // at a time, we only detect those if the type definition and the reflect usage
diff --git a/testdata/scripts/basic.txt b/testdata/scripts/basic.txt
index 3f42075..e70d43e 100644
--- a/testdata/scripts/basic.txt
+++ b/testdata/scripts/basic.txt
@@ -27,7 +27,7 @@ stdout 'unknown'
 ! stdout $gofullversion
 
 # The binary can't contain the version string either.
-! binsubstr main$exe ${WORK@R} 'globalVar' 'globalFunc' $gofullversion
+! binsubstr main$exe ${WORK@R} 'globalVar' 'globalFunc' $gofullversion 'z0.go'
 
 [short] stop # checking that the build is reproducible is slow