main.cpp

#include "BinaryMetadataExtractor.h"
#include "FilesystemMetadataExtractor.h"
#include <google/protobuf/descriptor.pb.h>
#include <google/protobuf/dynamic_message.h>
#include <google/protobuf/io/coded_stream.h>
#include <boost/filesystem.hpp>
#include <boost/program_options.hpp>
#include <fstream>
#include <iostream>
#include <list>
#include <unordered_map>

bool ParseFromCodedInputStreamWithDescriptorPool(google::protobuf::FileDescriptorProto& proto, google::protobuf::io::CodedInputStream* decoder)
{
    decoder->SetExtensionRegistry(google::protobuf::DescriptorPool::generated_pool(), google::protobuf::MessageFactory::generated_factory());
    return proto.ParseFromCodedStream(decoder) && decoder->ConsumedEntireMessage();
}

std::unordered_map<std::string, std::pair<google::protobuf::FileDescriptor const*, bool>> fileDescriptorsByName;
google::protobuf::DynamicMessageFactory* dynamicMessageFactory = new google::protobuf::DynamicMessageFactory(); // LEAKED ON PURPOSE, messages generated by this are double registered (and would be double freed)

int main(int argc, char* argv[])
{
    namespace po = boost::program_options;

    // Define command line options
    po::options_description desc("protobuf-decompiler: Reconstructs .proto files from protobuf descriptors\n\nOptions");
    desc.add_options()
        ("help,h", "Show this help message")
        ("binary", po::value<std::string>(), "Path to binary file containing protobuf descriptors")
        ("directory", po::value<std::string>()->default_value("."), "Directory to scan for .protoc files (when no binary is specified)")
        ("output,o", po::value<std::string>()->default_value("."), "Output directory for extracted .proto files")
    ;

    // Positional argument for binary path (for backward compatibility)
    po::positional_options_description p;
    p.add("binary", 1);

    po::variables_map vm;
    try {
        po::store(po::command_line_parser(argc, argv).
                  options(desc).positional(p).run(), vm);
        po::notify(vm);
    }
    catch (const po::error& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        std::cerr << "\nUse --help for usage information." << std::endl;
        return 1;
    }

    // Show help
    if (vm.count("help")) {
        std::cout << desc << "\n\n";
        std::cout << "Usage:\n";
        std::cout << "  " << argv[0] << " [options] [binary_file]\n";
        std::cout << "  " << argv[0] << " --binary <path_to_binary> [--output <output_dir>]\n";
        std::cout << "  " << argv[0] << " --directory <path_to_directory> [--output <output_dir>]\n\n";
        std::cout << "Examples:\n";
        std::cout << "  " << argv[0] << " /path/to/binary                      # Extract to current directory\n";
        std::cout << "  " << argv[0] << " /path/to/binary -o /path/to/output   # Extract to specific directory\n";
        std::cout << "  " << argv[0] << "                                      # Scan current directory for .protoc files\n";
        std::cout << "  " << argv[0] << " --help                               # Show this help\n";
        return 0;
    }

    // initialize Descriptor descriptors
    google::protobuf::DescriptorProto().GetMetadata();

    // first collect files
    std::unique_ptr<MetadataExtractor> extractor;
    if (vm.count("binary"))
    {
        // parse protobuf metadata from binary
        std::string binaryPath = vm["binary"].as<std::string>();
        extractor = std::make_unique<BinaryMetadataExtractor>();
        try {
            extractor->Parse(binaryPath);
        }
        catch (const std::exception& e) {
            std::cerr << "Error: " << e.what() << std::endl;
            return 1;
        }
    }
    else
    {
        std::string directory = vm["directory"].as<std::string>();
        extractor = std::make_unique<FilesystemMetadataExtractor>();
        extractor->Parse(boost::filesystem::path(directory));
    }

    if (extractor->GetMetadata().empty())
    {
        if (vm.count("binary"))
        {
            std::cout << "No .proto descriptors found in binary: " << vm["binary"].as<std::string>() << std::endl;
        }
        else
        {
            std::string directory = vm["directory"].as<std::string>();
            std::cout << "No .protoc files found in directory: " << boost::filesystem::absolute(directory) << std::endl;
            std::cout << std::endl;
            std::cout << "To extract from a binary, use: " << argv[0] << " --binary <path_to_binary>" << std::endl;
        }
        return 1;
    }

    // its time to hack private members
    google::protobuf::DescriptorPool* pool = google::protobuf::DescriptorPool::internal_generated_pool();
    void* mutex_ = *reinterpret_cast<char**>(pool);
    void* fallback_database_ = *(reinterpret_cast<void**>(pool) + 1);
    *reinterpret_cast<void**>(pool) = nullptr; // mutex_
    *(reinterpret_cast<void**>(pool) + 1) = nullptr; // fallback_database_

    std::vector<std::pair<google::protobuf::FileDescriptorProto, bool>> descProtos;
    std::set<std::string> parsed;
    std::list<std::string> parsedSorted;
    // then load everything into descriptor pool
    for (size_t i = 0; i < extractor->GetMetadata().size(); ++i)
    {
        for (std::unique_ptr<MetadataExtractor::Metadata> const& meta : extractor->GetMetadata())
        {
            if (parsed.count(meta->GetId()))
                continue;

            std::shared_ptr<google::protobuf::io::CodedInputStream> in = meta->CreateCodedInputStream();
            if (!in)
                continue;

            google::protobuf::FileDescriptorProto fileDescProto;
            bool parseOk = false;
            try
            {
                parseOk = ParseFromCodedInputStreamWithDescriptorPool(fileDescProto, in.get());
            }
            catch (...)
            {
                continue;
            }

            if (parseOk)
            {
                if (google::protobuf::FileDescriptor const* fileDesc = pool->BuildFile(fileDescProto))
                {
                    if (fileDesc->name() != "google/protobuf/descriptor.proto" && fileDescriptorsByName.count(fileDesc->name()) == 0)
                    {
                        auto itr = fileDescriptorsByName.emplace(std::piecewise_construct, std::forward_as_tuple(fileDesc->name()), std::forward_as_tuple(fileDesc, false)).first;
                        google::protobuf::MessageFactory::InternalRegisterGeneratedFile(itr->first.c_str(), [](std::string const& name)
                        {
                            std::pair<google::protobuf::FileDescriptor const*, bool>& p = fileDescriptorsByName[name];
                            if (p.second)
                                return;

                            p.second = true;
                            google::protobuf::FileDescriptor const* desc = p.first;
                            for (int j = 0; j < desc->message_type_count(); ++j)
                                google::protobuf::MessageFactory::InternalRegisterGeneratedMessage(desc->message_type(j), dynamicMessageFactory->GetPrototype(desc->message_type(j)));
                        });
                    }

                    for (int j = 0; j < fileDesc->extension_count(); ++j)
                    {
                        google::protobuf::FieldDescriptor const* extension = fileDesc->extension(j);
                        switch (extension->type())
                        {
                            case google::protobuf::FieldDescriptor::Type::TYPE_ENUM:
                                google::protobuf::internal::ExtensionSet::RegisterEnumExtension(
                                    google::protobuf::MessageFactory::generated_factory()->GetPrototype(extension->containing_type()),
                                    extension->number(),
                                    extension->type(),
                                    extension->is_repeated(),
                                    extension->is_packed(),
                                    [](int){ return true; });
                                break;
                            case google::protobuf::FieldDescriptor::Type::TYPE_MESSAGE:
                            case google::protobuf::FieldDescriptor::Type::TYPE_GROUP:
                                google::protobuf::internal::ExtensionSet::RegisterMessageExtension(
                                    google::protobuf::MessageFactory::generated_factory()->GetPrototype(extension->containing_type()),
                                    extension->number(),
                                    extension->type(),
                                    extension->is_repeated(),
                                    extension->is_packed(),
                                    google::protobuf::MessageFactory::generated_factory()->GetPrototype(extension->message_type()));
                                break;
                            default:
                                google::protobuf::internal::ExtensionSet::RegisterExtension(
                                    google::protobuf::MessageFactory::generated_factory()->GetPrototype(extension->containing_type()),
                                    extension->number(),
                                    extension->type(),
                                    extension->is_repeated(),
                                    extension->is_packed());
                                break;
                        }
                    }

                    parsed.insert(meta->GetId());
                    if (fileDesc->name() != "google/protobuf/descriptor.proto")
                        parsedSorted.push_back(meta->GetId());
                    else
                        parsedSorted.push_front(meta->GetId());
                }
            }
        }
    }

    // Get output directory
    boost::filesystem::path outputDir(vm["output"].as<std::string>());
    try {
        if (!boost::filesystem::exists(outputDir))
        {
            boost::filesystem::create_directories(outputDir);
        }
        else if (!boost::filesystem::is_directory(outputDir))
        {
            std::cerr << "Error: Output path exists but is not a directory: " << outputDir << std::endl;
            return 1;
        }
        
        // Test if directory is writable by creating a temporary file
        boost::filesystem::path testFile = outputDir / ".protobuf_decompiler_test";
        std::ofstream test(testFile.string());
        if (!test.good())
        {
            std::cerr << "Error: Output directory is not writable: " << outputDir << std::endl;
            return 1;
        }
        test.close();
        boost::filesystem::remove(testFile);
    }
    catch (const boost::filesystem::filesystem_error& e)
    {
        std::cerr << "Error: Failed to create output directory: " << e.what() << std::endl;
        return 1;
    }

    // and finally rebuild all protos in a new pool with fully resolved dependencies and extensions
    google::protobuf::DescriptorPool* pool2 = new google::protobuf::DescriptorPool();
    for (std::string const& fileName : parsedSorted)
    {
        MetadataExtractor::Metadata const* meta = extractor->GetById(fileName);
        if (!meta)
            continue;

        std::shared_ptr<google::protobuf::io::CodedInputStream> in = meta->CreateCodedInputStream();
        if (!in)
            continue;

        google::protobuf::FileDescriptorProto fileDescProto;
        if (ParseFromCodedInputStreamWithDescriptorPool(fileDescProto, in.get()))
        {
            fileDescProto.mutable_options()->set_optimize_for(google::protobuf::FileOptions_OptimizeMode_SPEED);
            if (google::protobuf::FileDescriptor const* fileDesc = pool2->BuildFile(fileDescProto))
            {
                boost::filesystem::path outputPath = outputDir / fileDesc->name();
                boost::filesystem::path parentPath = outputPath.parent_path();
                if (!parentPath.empty())
                    boost::filesystem::create_directories(parentPath);
                std::ofstream f(outputPath.string());
                f << fileDesc->DebugString() << std::endl;
                f.close();
            }
        }
    }

    // unhack to free memory
    *reinterpret_cast<void**>(pool) = mutex_;
    *(reinterpret_cast<void**>(pool) + 1) = fallback_database_;

    delete pool2;

    google::protobuf::ShutdownProtobufLibrary();

    return 0;
}