func (p *parser) parseImportSpec(n *parse.Node) *ast.ImportSpec {
spec := ast.ImportSpec{}
if n.Child(0).Is(importPath) {
spec.Path = p.parseBasicLit(n.Child(0))
} else {
name := n.Child(0)
switch name.Rule() {
case identifier:
spec.Name = p.parseIdent(name)
case term("."):
spec.Name = &ast.Ident{
NamePos: token.Pos(name.Pos()),
Name: ".",
}
}
spec.Path = p.parseBasicLit(n.Child(1))
}
return &spec
}
func fixGoExact(f *ast.File) bool {
// This one is harder because the import name changes.
// First find the import spec.
var importSpec *ast.ImportSpec
walk(f, func(n interface{}) {
if importSpec != nil {
return
}
spec, ok := n.(*ast.ImportSpec)
if !ok {
return
}
path, err := strconv.Unquote(spec.Path.Value)
if err != nil {
return
}
if path == "golang.org/x/tools/go/exact" {
importSpec = spec
}
})
if importSpec == nil {
return false
}
// We are about to rename exact.* to constant.*, but constant is a common
// name. See if it will conflict. This is a hack but it is effective.
exists := renameTop(f, "constant", "constant")
suffix := ""
if exists {
suffix = "_"
}
// Now we need to rename all the uses of the import. RewriteImport
// affects renameTop, but not vice versa, so do them in this order.
renameTop(f, "exact", "constant"+suffix)
rewriteImport(f, "golang.org/x/tools/go/exact", "go/constant")
// renameTop will also rewrite the imported package name. Fix that;
// we know it should be missing.
importSpec.Name = nil
return true
}
func main() {
flag.Parse()
src, err := ioutil.ReadFile(*in)
if err != nil {
log.Fatal(err)
}
fset := token.NewFileSet()
f, _ := parser.ParseFile(fset, "", src, parser.AllErrors)
// Find where "crypto/aes" is imported.
var importSpec *ast.ImportSpec
for _, imp := range f.Imports {
if imp.Path.Value == target {
importSpec = imp
break
}
}
if importSpec == nil {
log.Fatalf("File doesn't import %v!", target)
}
// Name the import explicitly if it's not already named. Change the import path to a white-box construction.
if importSpec.Name == nil {
importSpec.Name = &ast.Ident{0, "aes", nil}
} else if importSpec.Name.Name == "." {
log.Fatalf("Can't transform a file's encryption keys if %v is imported into the local scope!", target)
}
importSpec.Path.Value = replacement
// Find everywhere we initialize a block cipher with an explicit key and replace it with a white-box.
ast.Inspect(f, func(n ast.Node) bool {
switch n.(type) {
case *ast.CallExpr:
callExpr := n.(*ast.CallExpr)
if IsCallToEncrypt(importSpec.Name.Name, callExpr) {
key, ok := ExtractKey(callExpr.Args[0]) // aes.NewCipher has only one argument.
if ok {
TransformCallToEncrypt(callExpr, key)
return false
} else {
log.Printf("Found encryption call at %v, but couldn't extract the key!", fset.Position(n.Pos()))
return true
}
}
}
return true
})
dst, err := os.Create(*out)
if err != nil {
log.Fatal(err)
}
printer.Fprint(dst, fset, f)
}