// SPDX-License-Identifier: MIT
// apps/src/dso-sign.cpp
//
// dso-sign — sign a DSO (shared library) and embed manifest as ELF sections
//
// Usage:
//   dso-sign addon.so --key developer.key
//   dso-sign addon.so --key developer.key --name my-addon --version 1.2.0
//   dso-sign addon.so --key developer.key --author "Name <email>"

#include <sign/Manifest.hh>

#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/err.h>

#include <chrono>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <sstream>
#include <string>
#include <vector>

namespace fs = std::filesystem;
using namespace fedem::sign;

static void fatalSsl( const char* msg )
{
    std::fprintf( stderr, "error: %s\n", msg );
    ERR_print_errors_fp( stderr );
    std::exit( 1 );
}

static std::string hexEncode( const unsigned char* data, std::size_t len )
{
    std::ostringstream ss;
    ss << std::hex << std::setfill('0');
    for( std::size_t i = 0; i < len; ++i )
        ss << std::setw(2) << static_cast<unsigned>(data[i]);
    return ss.str();
}

static std::string sha256File( std::string const& path )
{
    std::ifstream f( path, std::ios::binary );
    if( !f ) { std::fprintf( stderr, "error: cannot open %s\n", path.c_str() ); std::exit(1); }
    EVP_MD_CTX* ctx = EVP_MD_CTX_new();
    EVP_DigestInit_ex( ctx, EVP_sha256(), nullptr );
    char buf[65536];
    while( f.read(buf, sizeof buf) || f.gcount() > 0 )
        EVP_DigestUpdate( ctx, buf, static_cast<std::size_t>(f.gcount()) );
    unsigned char digest[EVP_MAX_MD_SIZE]; unsigned int dlen = 0;
    EVP_DigestFinal_ex( ctx, digest, &dlen );
    EVP_MD_CTX_free( ctx );
    return hexEncode( digest, dlen );
}

static std::string isoTimestamp()
{
    auto t = std::chrono::system_clock::to_time_t( std::chrono::system_clock::now() );
    std::tm* tm = std::gmtime( &t );
    char buf[32]; std::strftime( buf, sizeof buf, "%Y-%m-%dT%H:%M:%SZ", tm );
    return buf;
}

static std::string pubKeyFingerprint( EVP_PKEY* pkey )
{
    unsigned char* der = nullptr; int len = i2d_PUBKEY( pkey, &der );
    if( len <= 0 ) fatalSsl( "i2d_PUBKEY" );
    std::size_t dlen = 0; unsigned char digest[EVP_MAX_MD_SIZE];
    EVP_Q_digest( nullptr, "SHA256", nullptr, der, static_cast<std::size_t>(len), digest, &dlen );
    OPENSSL_free( der );
    return hexEncode( digest, dlen );
}

static void writeTempFile( std::string const& path, const void* data, std::size_t len )
{
    FILE* f = std::fopen( path.c_str(), "wb" );
    if( !f ) { std::fprintf( stderr, "error: cannot write %s\n", path.c_str() ); std::exit(1); }
    std::fwrite( data, 1, len, f );
    std::fclose( f );
}

static void stripExistingSections( std::string const& path )
{
    std::string cmd =
        "objcopy --remove-section " + std::string(kManifestSection) +
        " --remove-section " + std::string(kSigSection) +
        " \"" + path + "\" 2>/dev/null";
    std::system( cmd.c_str() );
}

int main( int argc, char** argv )
{
    std::string dsoPath, keyPath, addonName, addonVersion, authorOverride;
    int abiMajor = 1, abiMinor = 0;
    bool quiet = false;

    for( int i = 1; i < argc; ++i )
    {
        if(      !std::strcmp(argv[i],"--key")     && i+1<argc ) keyPath        = argv[++i];
        else if( !std::strcmp(argv[i],"--name")    && i+1<argc ) addonName      = argv[++i];
        else if( !std::strcmp(argv[i],"--version") && i+1<argc ) addonVersion   = argv[++i];
        else if( !std::strcmp(argv[i],"--author")  && i+1<argc ) authorOverride = argv[++i];
        else if( !std::strcmp(argv[i],"--abi-major")&& i+1<argc ) abiMajor      = std::atoi(argv[++i]);
        else if( !std::strcmp(argv[i],"--abi-minor")&& i+1<argc ) abiMinor      = std::atoi(argv[++i]);
        else if( !std::strcmp(argv[i],"--quiet") || !std::strcmp(argv[i],"-q") ) quiet = true;
        else if( !std::strcmp(argv[i],"--help") || !std::strcmp(argv[i],"-h") )
        {
            std::puts("Usage: dso-sign <addon.so> --key <developer.key>\n"
                      "       [--name <n>] [--version <v>] [--author \"Name <email>\"]\n"
                      "       [--abi-major N] [--abi-minor N] [--quiet]");
            return 0;
        }
        else if( argv[i][0] != '-' ) dsoPath = argv[i];
        else { std::fprintf(stderr,"error: unknown argument: %s\n",argv[i]); return 1; }
    }

    if( dsoPath.empty() || keyPath.empty() )
    { std::fputs("error: DSO path and --key are required\n",stderr); return 1; }
    if( !fs::exists(dsoPath) ) { std::fprintf(stderr,"error: not found: %s\n",dsoPath.c_str()); return 1; }
    if( !fs::exists(keyPath) ) { std::fprintf(stderr,"error: key not found: %s\n",keyPath.c_str()); return 1; }

    if( addonName.empty() )    addonName    = fs::path(dsoPath).stem().string();
    if( addonVersion.empty() ) addonVersion = "1.0.0";

    if( !quiet ) std::printf("[dso-sign] stripping existing sections...\n");
    stripExistingSections( dsoPath );

    if( !quiet ) std::printf("[dso-sign] computing SHA-256...\n");
    std::string const fileHash = sha256File( dsoPath );
    if( !quiet ) std::printf("[dso-sign] sha256: %s\n", fileHash.c_str() );

    FILE* kf = std::fopen( keyPath.c_str(), "r" );
    if( !kf ) { std::fprintf(stderr,"error: cannot open %s\n",keyPath.c_str()); return 1; }
    EVP_PKEY* pkey = PEM_read_PrivateKey( kf, nullptr, nullptr, nullptr );
    std::fclose( kf );
    if( !pkey ) fatalSsl("PEM_read_PrivateKey");

    // Author from key comment
    std::string author = authorOverride;
    if( author.empty() )
    {
        std::ifstream kfile( keyPath );
        std::string line, kname, kemail;
        while( std::getline(kfile,line) )
        {
            if( line.rfind("# Name:",0)==0 )  kname  = line.substr(8);
            if( line.rfind("# Email:",0)==0 ) kemail = line.substr(9);
        }
        auto trim=[](std::string& s){ while(!s.empty()&&s.front()==' ')s.erase(0,1); while(!s.empty()&&s.back()==' ')s.pop_back(); };
        trim(kname); trim(kemail);
        if( !kname.empty() && !kemail.empty() ) author = kname + " <" + kemail + ">";
        else if( !kname.empty() ) author = kname;
    }

    std::string const keyId = pubKeyFingerprint( pkey );
    if( !quiet ) std::printf("[dso-sign] key fingerprint: %s\n", keyId.c_str());

    ManifestData md;
    md.name        = addonName;
    md.version     = addonVersion;
    md.author      = author;
    md.sha256      = fileHash;
    md.publicKeyId = keyId;
    md.abiMajor    = abiMajor;
    md.abiMinor    = abiMinor;
    md.timestamp   = isoTimestamp();

    std::string const manifestJson = Manifest::toJson( md );
    if( !quiet ) std::printf("[dso-sign] manifest:\n%s\n", manifestJson.c_str());

    // Sign manifest
    EVP_MD_CTX* mctx = EVP_MD_CTX_new();
    if( EVP_DigestSignInit(mctx,nullptr,nullptr,nullptr,pkey) <= 0 ) fatalSsl("EVP_DigestSignInit");
    auto const* msg = reinterpret_cast<const unsigned char*>(manifestJson.data());
    std::size_t sigLen = 0;
    if( EVP_DigestSign(mctx,nullptr,&sigLen,msg,manifestJson.size()) <= 0 ) fatalSsl("EVP_DigestSign size");
    std::vector<unsigned char> sig(sigLen);
    if( EVP_DigestSign(mctx,sig.data(),&sigLen,msg,manifestJson.size()) <= 0 ) fatalSsl("EVP_DigestSign");
    EVP_MD_CTX_free( mctx );
    EVP_PKEY_free( pkey );

    // Write temp files + objcopy
    std::string const tmpManifest = dsoPath + ".manifest.tmp";
    std::string const tmpSig      = dsoPath + ".sig.tmp";
    writeTempFile( tmpManifest, manifestJson.data(), manifestJson.size() );
    writeTempFile( tmpSig,      sig.data(),           sigLen );

    std::string cmd =
        "objcopy"
        " --add-section " + std::string(kManifestSection) + "=" + tmpManifest +
        " --set-section-flags " + std::string(kManifestSection) + "=noload,readonly" +
        " --add-section " + std::string(kSigSection) + "=" + tmpSig +
        " --set-section-flags " + std::string(kSigSection) + "=noload,readonly" +
        " \"" + dsoPath + "\"";

    if( !quiet ) std::printf("[dso-sign] embedding sections...\n");
    if( std::system(cmd.c_str()) != 0 )
    { std::fputs("error: objcopy failed\n",stderr); fs::remove(tmpManifest); fs::remove(tmpSig); return 1; }

    fs::remove( tmpManifest );
    fs::remove( tmpSig );

    std::printf("[dso-sign] done: %s is signed\n", dsoPath.c_str());
    return 0;
}