From 8df134ab792511a04f5153cec457966919aebbbb Mon Sep 17 00:00:00 2001 From: Arkadiy Tetelman Date: Thu, 29 Dec 2022 16:52:51 -0800 Subject: [PATCH] Initial commit --- .gitignore | 1 + Makefile | 5 + README.md | 0 cmd/main.go | 63 ++++++++++ go.mod | 15 +++ go.sum | 26 ++++ pkg/proto/proto.go | 260 ++++++++++++++++++++++++++++++++++++++++ pkg/proto/proto_test.go | 18 +++ pkg/scan/scan.go | 58 +++++++++ pkg/scan/scan_test.go | 1 + 10 files changed, 447 insertions(+) create mode 100644 .gitignore create mode 100644 Makefile create mode 100644 README.md create mode 100644 cmd/main.go create mode 100644 go.mod create mode 100644 go.sum create mode 100644 pkg/proto/proto.go create mode 100644 pkg/proto/proto_test.go create mode 100644 pkg/scan/scan.go create mode 100644 pkg/scan/scan_test.go diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..91c1f53 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +protodump diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..755bd29 --- /dev/null +++ b/Makefile @@ -0,0 +1,5 @@ +build: + go build -o protodump cmd/main.go + +test: + go test -v ./... diff --git a/README.md b/README.md new file mode 100644 index 0000000..e69de29 diff --git a/cmd/main.go b/cmd/main.go new file mode 100644 index 0000000..d6a21eb --- /dev/null +++ b/cmd/main.go @@ -0,0 +1,63 @@ +package main + +import ( + "flag" + "fmt" + "log" + "os" + "path" + "strings" + + "github.com/arkadiyt/protodump/pkg/proto" + "github.com/arkadiyt/protodump/pkg/scan" +) + +var debug bool + +func Debug(str string, a ...any) (int, error) { + if debug { + return fmt.Printf(str, a...) + } + return 0, nil +} + +func main() { + cwd, err := os.Getwd() + if err != nil { + log.Fatalf("Couldn't determine current working directory: %v\n", err) + } + + var file = flag.String("file", "", "The file to extract definitions from") + var output = flag.String("output", cwd, "The output directory to save definitions in (will be created if it doesn't exist). Defaults to current directory.") + flag.BoolVar(&debug, "v", false, "Verbose output") + flag.Parse() + + if *file == "" { + fmt.Printf("Usage:\n") + flag.PrintDefaults() + return + } + + results, err := scan.ScanFile(*file) + if err != nil { + log.Fatalf("Got error scanning: %v\n", err) + } + + for _, result := range results { + definition, err := proto.NewFromBytes(result) + if err != nil { + Debug("Got error parsing definition: %v\n", err) + } else { + filename := definition.Filename() + if strings.HasSuffix(filename, ".proto") { + dir := path.Join(*output, path.Dir(filename)) + final := path.Join(dir, path.Base(filename)) + os.MkdirAll(dir, 0700) + os.WriteFile(final, []byte(definition.String()), 0700) + fmt.Printf("Wrote %s\n", final) + } else { + // Need to investigate further + } + } + } +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..e6a3730 --- /dev/null +++ b/go.mod @@ -0,0 +1,15 @@ +module github.com/arkadiyt/protodump + +go 1.19 + +require ( + golang.org/x/exp v0.0.0-20221229185413-b6b61370dfa8 + google.golang.org/protobuf v1.28.1 +) + +require ( + github.com/davecgh/go-spew v1.1.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + github.com/stretchr/testify v1.8.1 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..cad774a --- /dev/null +++ b/go.sum @@ -0,0 +1,26 @@ +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= +github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk= +github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= +golang.org/x/exp v0.0.0-20221227203929-1b447090c38c h1:Govq2W3bnHJimHT2ium65kXcI7ZzTniZHcFATnLJM0Q= +golang.org/x/exp v0.0.0-20221227203929-1b447090c38c/go.mod h1:CxIveKay+FTh1D0yPZemJVgC/95VzuuOLq5Qi4xnoYc= +golang.org/x/exp v0.0.0-20221229185413-b6b61370dfa8 h1:/09aNBEyXdSQMyJEalp+UtcHlmXZ+q7nTJAQppUpkV4= +golang.org/x/exp v0.0.0-20221229185413-b6b61370dfa8/go.mod h1:CxIveKay+FTh1D0yPZemJVgC/95VzuuOLq5Qi4xnoYc= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= +google.golang.org/protobuf v1.28.1 h1:d0NfwRgPtno5B1Wa6L2DAG+KivqkdutMf1UhdNx175w= +google.golang.org/protobuf v1.28.1/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/pkg/proto/proto.go b/pkg/proto/proto.go new file mode 100644 index 0000000..17584a7 --- /dev/null +++ b/pkg/proto/proto.go @@ -0,0 +1,260 @@ +package proto + +import ( + "fmt" + "strconv" + "strings" + + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/reflect/protodesc" + "google.golang.org/protobuf/reflect/protoreflect" + "google.golang.org/protobuf/reflect/protoregistry" + "google.golang.org/protobuf/types/descriptorpb" +) + +type ProtoDefinition struct { + builder strings.Builder + indendation int + pb descriptorpb.FileDescriptorProto + descriptor protoreflect.FileDescriptor + filename string +} + +// TODO add proto2 support + +func (pd *ProtoDefinition) indent() { + pd.indendation += 1 +} + +func (pd *ProtoDefinition) dedent() { + pd.indendation -= 1 +} + +func (pd *ProtoDefinition) writeIndented(s string) { + pd.builder.WriteString(strings.Repeat(" ", pd.indendation)) + pd.write(s) +} + +func (pd *ProtoDefinition) write(s string) { + pd.builder.WriteString(s) +} + +func (pd *ProtoDefinition) String() string { + return pd.builder.String() +} + +func (pd *ProtoDefinition) Filename() string { + return pd.filename +} + +func (pd *ProtoDefinition) writeMethod(method protoreflect.MethodDescriptor) { + // TODO need to handle method options + pd.writeIndented("rpc ") + pd.write(string(method.Name())) + pd.write("(") + if method.IsStreamingClient() { + pd.write("streaming ") + } + pd.write(string(method.Input().Name())) + pd.write(") returns (") + if method.IsStreamingServer() { + pd.write("streaming ") + } + pd.write(string(method.Output().Name())) + pd.write(") {}\n") +} + +func (pd *ProtoDefinition) writeService(service protoreflect.ServiceDescriptor) { + // TODO need to handle service options + pd.write("service ") + pd.write(string(service.Name())) + pd.write(" {\n") + pd.indent() + for i := 0; i < service.Methods().Len(); i++ { + pd.writeMethod(service.Methods().Get(i)) + } + pd.dedent() + pd.writeIndented("}\n\n") +} + +func (pd *ProtoDefinition) writeType(field protoreflect.FieldDescriptor) { + kind := field.Kind().String() + if kind == "message" { + pd.write(string(field.Message().Name())) + } else if kind == "map" { + pd.write("map<") + pd.writeType(field.MapKey()) + pd.write(", ") + pd.writeType(field.MapValue()) + pd.write(">") + } else { + pd.write(kind) + } +} + +func (pd *ProtoDefinition) writeOneof(oneof protoreflect.OneofDescriptor) { + // TODO need to handle oneof options + pd.writeIndented("") + pd.write("oneof ") + pd.write(string(oneof.Name())) + pd.write(" {\n") + pd.indent() + for i := 0; i < oneof.Fields().Len(); i++ { + pd.writeField(oneof.Fields().Get(i)) + } + pd.dedent() + pd.writeIndented("}\n") +} + +func (pd *ProtoDefinition) writeField(field protoreflect.FieldDescriptor) { + // TODO need to handle options + pd.writeIndented("") + if field.HasOptionalKeyword() { + pd.write("optional ") + } else if field.Cardinality().String() == "repeated" { + pd.write("repeated ") + } + pd.writeType(field) + pd.write(" ") + pd.write(string(field.Name())) + pd.write(" = ") + pd.write(strconv.Itoa(field.Index())) + pd.write(";\n") +} + +func (pd *ProtoDefinition) writeEnum(enum protoreflect.EnumDescriptor) { + pd.write("enum ") + pd.write(string(enum.Name())) + pd.write(" {\n") + // TODO need to handle enum options (allow_alias) + pd.indent() + for i := 0; i < enum.Values().Len(); i++ { + value := enum.Values().Get(i) + pd.writeIndented(string(value.Name())) + pd.write(" = ") + pd.write(fmt.Sprintf("%d", value.Number())) + pd.write(";\n") + } + pd.dedent() + pd.writeIndented("}\n\n") +} + +func (pd *ProtoDefinition) writeMessage(message protoreflect.MessageDescriptor) { + // TODO need to handle message options + pd.write("message ") + pd.write(string(message.Name())) + pd.write(" {\n") + pd.indent() + + for i := 0; i < message.ReservedNames().Len(); i++ { + name := message.ReservedNames().Get(i) + pd.writeIndented("reserved ") + pd.write(string(name)) + pd.write(";\n") + } + + for i := 0; i < message.ReservedRanges().Len(); i++ { + reservedRange := message.ReservedRanges().Get(i) + pd.writeIndented("reserved ") + if reservedRange[0] == reservedRange[1] { + pd.write(string(reservedRange[0])) + } else { + if reservedRange[0] > reservedRange[1] { + reservedRange[1], reservedRange[0] = reservedRange[0], reservedRange[1] + } + pd.write(fmt.Sprintf("%d", reservedRange[0])) + pd.write(" to ") + pd.write(fmt.Sprintf("%d", reservedRange[1])) + } + pd.write(";\n") + } + + for i := 0; i < message.Fields().Len(); i++ { + field := message.Fields().Get(i) + if field.ContainingOneof() == nil { + pd.writeField(field) + } + } + + for i := 0; i < message.Oneofs().Len(); i++ { + pd.writeOneof(message.Oneofs().Get(i)) + } + pd.dedent() + pd.writeIndented("}\n\n") +} + +func (pd *ProtoDefinition) writeImport(fileImport protoreflect.FileImport) { + pd.write("import ") + if fileImport.IsPublic { + pd.write("public ") + } + pd.write("\"") + pd.write(fileImport.Path()) + pd.write("\";\n") +} + +func (pd *ProtoDefinition) writeFileDescriptor() { + pd.write("syntax = \"") + pd.write(pd.descriptor.Syntax().String()) + pd.write("\"\n\n") + + pd.write("package ") + pd.write(string(pd.descriptor.Package().Name())) + pd.write(";\n\n") + + // TODO need to handle FileOptions + filepb := protodesc.ToFileDescriptorProto(pd.descriptor) + options := filepb.Options + if options != nil { + } + + for i := 0; i < pd.descriptor.Imports().Len(); i++ { + pd.writeImport(pd.descriptor.Imports().Get(i)) + } + + if pd.descriptor.Imports().Len() > 0 { + pd.write("\n") + } + + for i := 0; i < pd.descriptor.Services().Len(); i++ { + pd.writeService(pd.descriptor.Services().Get(i)) + } + + for i := 0; i < pd.descriptor.Messages().Len(); i++ { + pd.writeMessage(pd.descriptor.Messages().Get(i)) + } + + for i := 0; i < pd.descriptor.Enums().Len(); i++ { + pd.writeEnum(pd.descriptor.Enums().Get(i)) + } +} + +func NewFromBytes(payload []byte) (*ProtoDefinition, error) { + var pb descriptorpb.FileDescriptorProto + err := proto.Unmarshal(payload, &pb) + if err != nil { + return nil, fmt.Errorf("Couldn't unmarshal proto: %w", err) + } + + return NewFromDescriptor(pb) +} + +func NewFromDescriptor(pb descriptorpb.FileDescriptorProto) (*ProtoDefinition, error) { + fileOptions := protodesc.FileOptions{AllowUnresolvable: true} + descriptor, err := fileOptions.New(&pb, &protoregistry.Files{}) + + if err != nil { + return nil, fmt.Errorf("Couldn't create FileDescriptor: %w", err) + } + + pd := ProtoDefinition{ + pb: pb, + descriptor: descriptor, + filename: descriptor.Path(), + } + + pd.writeFileDescriptor() + + return &pd, nil + +} diff --git a/pkg/proto/proto_test.go b/pkg/proto/proto_test.go new file mode 100644 index 0000000..fe7a47c --- /dev/null +++ b/pkg/proto/proto_test.go @@ -0,0 +1,18 @@ +package proto + +import ( + "os" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestDefinitions(t *testing.T) { + files, err := os.ReadDir("fixtures") + assert.NoError(t, err) + + for _, file := range files { + t.Run(file.Name(), func(t *testing.T) { + }) + } +} diff --git a/pkg/scan/scan.go b/pkg/scan/scan.go new file mode 100644 index 0000000..cdc2d9c --- /dev/null +++ b/pkg/scan/scan.go @@ -0,0 +1,58 @@ +package scan + +import ( + "bytes" + "fmt" + "os" + + "google.golang.org/protobuf/encoding/protowire" +) + +func ScanFile(path string) ([][]byte, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("Couldn't open file: %w", err) + } + return Scan(data), nil +} + +func consumeBytes(data []byte, position int) int { + start := position + for { + _, _, len := protowire.ConsumeField(data[position:]) + if len < 0 { + _ = protowire.ParseError(len) + return position - start + + } + position += len + } +} + +func Scan(data []byte) [][]byte { + results := make([][]byte, 0) + + for { + index := bytes.Index(data, []byte(".proto")) + if index == -1 { + break + } + + // Look for a 0xa byte: 0 0001 010 + // ^ no more bytes for varint + // ^ field 1 (file name) + // ^ type 2 (LEN) + start := bytes.LastIndexByte(data[:index], 0xa) + if start == -1 { + // fmt.Printf("Warn: couldn't find start to proto definition at offset:0x%s\n", strconv.FormatInt(int64(index), 16)) + data = data[index+1:] + continue + } + + length := consumeBytes(data, start) + results = append(results, data[start:start+length]) + data = data[start+length:] + } + + return results +} diff --git a/pkg/scan/scan_test.go b/pkg/scan/scan_test.go new file mode 100644 index 0000000..040485b --- /dev/null +++ b/pkg/scan/scan_test.go @@ -0,0 +1 @@ +package scan