// compile with: cl /EHsc xda2nbftool.cpp advapi32.lib

#include <stdio.h>
#include <windows.h>
#include <wincrypt.h>

bool bVerbose= false;

#include <vector>
typedef std::vector<BYTE> ByteVector;
typedef std::vector<DWORD> DwordVector;
#define vectorptr(v)  (&(v)[0])

//--------------------------------------------------------------------
static class CRCCalc {
public:
    CRCCalc() {
        initcrc(0xEDB88320L);
    }

    DWORD crc32(DWORD prev, BYTE c) const
    {
       return (prev>>8) ^ crc32tab[(int)((prev&0xff)^c)];
    }

private:
        DWORD crc32tab[256];
 
    /* crcBase = 0xEDB88320L */
    DWORD calccrc32tab(BYTE c, DWORD crcBase)
    {  
       DWORD val = c;
       int i;
       for (i=0 ; i<8 ; i++)
       {
          if (val&1)
            val = (val>>1)^crcBase;
          else
            val = (val>>1);
       }
       return val;
    }

    void initcrc(DWORD crcBase)
    {
       DWORD val;
       int i;
       for (i=0 ; i<256 ; i++)
       {
          val = calccrc32tab(i, crcBase);
          crc32tab[i]=val;
       }
    }
} g_crccalc;

DWORD add_to_crc32(DWORD prev, BYTE c)
{
    return g_crccalc.crc32(prev, c);
}

DWORD calccrc32(const BYTE* buf, int buflen)
{
    DWORD crc= 0;

    while(buflen--) 
        crc= add_to_crc32(crc, *buf++);

    return crc;
}
//--------------------------------------------------------------------
void error(char *msg)
{
    printf("ERROR: %s\n", msg);
}
void dump_nbf_header(char *buf)
{
    char buf2[64];
    memcpy(buf2, buf, 64);
    buf2[10]= buf2[31]= buf2[42]= buf2[52]= buf2[61]= 0;
    printf("devicetype= %s\n", buf2);
    printf("operatorname= %s\n", buf2+11);
    printf("language= %s\n", buf2+32);
    printf("version= %s\n", buf2+43);
    printf("crc= %s\n", buf2+53);
}

bool encrypt_file(char *szDecryptedFilename, char *szEncryptedFilename, char *szKey)
{
    FILE *fin= fopen(szDecryptedFilename, "rb");
    if (fin==NULL) {
        perror(szDecryptedFilename);
        return false;
    }
    FILE *fout= fopen(szEncryptedFilename, "wb");
    if (fout==NULL) {
        perror(szEncryptedFilename);
        return false;
    }
    HCRYPTPROV hProv;
    if (!CryptAcquireContext(&hProv, NULL, "Microsoft Enhanced Cryptographic Provider v1.0", PROV_RSA_FULL, 0))
    {
        error("CryptAcquireContext");
        return false;
    }

    HCRYPTHASH hHash;
    if (!CryptCreateHash(hProv, CALG_MD5, NULL, NULL, &hHash))
    {
        error("CryptCreateHash");
        return false;
    }
    if (!CryptHashData(hHash, (BYTE*)szKey, strlen(szKey), 0))
    {
        error("CryptHashData");
        return false;
    }
    HCRYPTKEY hKey;
    if (!CryptDeriveKey(hProv, ALG_CLASS_DATA_ENCRYPT | ALG_TYPE_STREAM | ALG_SID_DES, hHash, 128<<16, &hKey))
    {
        error("CryptDeriveKey");
        return false;
    }
    if (!CryptDestroyHash(hHash))
    {
        error("CryptDestroyHash");
        return false;
    }
    hHash= 0;

    ByteVector buffer; buffer.resize(0x40000);

    DWORD crcsum= 0;
    DWORD filecrc= 0;

    bool bFlag= true;
    do {
        DWORD dwDataLen= fread(vectorptr(buffer), 1, bFlag?0x40:0x40000, fin);

        if (bFlag) {
            if (!(buffer[10]=='-' && buffer[31]=='-' && buffer[42]=='-' && buffer[52]=='-' && buffer[61]=='-'))
                printf("WARNING: this does not look like a nbf header\n");
            else
                filecrc=strtoul((char*)vectorptr(buffer)+0x35,0,16);
        }


        crcsum += calccrc32(vectorptr(buffer), bFlag?0x34:dwDataLen);

        if (!CryptEncrypt(hKey, NULL, feof(fin), NULL, vectorptr(buffer), &dwDataLen, dwDataLen))
        {
            error("CryptEncrypt");
            return false;
        }
        bFlag= false;

        fwrite(vectorptr(buffer), 1, dwDataLen, fout);
    } while (!feof(fin));

    fclose(fin);
    fclose(fout);
    CryptDestroyKey(hKey);
    CryptReleaseContext(hProv, 0);

    if (filecrc != crcsum)
        printf("WARNING: checksum is not ok\n");

    return true;
}

bool decrypt_file(char *szEncryptedFilename, char *szDecryptedFilename, char *szKey)
{
    FILE *fin= fopen(szEncryptedFilename, "rb");
    if (fin==NULL) {
        perror(szEncryptedFilename);
        return false;
    }
    FILE *fout= fopen(szDecryptedFilename, "wb");
    if (fout==NULL) {
        perror(szDecryptedFilename);
        return false;
    }
    HCRYPTPROV hProv;
    if (!CryptAcquireContext(&hProv, NULL, "Microsoft Enhanced Cryptographic Provider v1.0", PROV_RSA_FULL, 0))
    {
        error("CryptAcquireContext");
        return false;
    }

    HCRYPTHASH hHash;
    if (!CryptCreateHash(hProv, CALG_MD5, NULL, NULL, &hHash))
    {
        error("CryptCreateHash");
        return false;
    }
    if (!CryptHashData(hHash, (BYTE*)szKey, strlen(szKey), 0))
    {
        error("CryptHashData");
        return false;
    }
    HCRYPTKEY hKey;
    if (!CryptDeriveKey(hProv, ALG_CLASS_DATA_ENCRYPT | ALG_TYPE_STREAM | ALG_SID_DES, hHash, 128<<16, &hKey))
    {
        error("CryptDeriveKey");
        return false;
    }
    if (!CryptDestroyHash(hHash))
    {
        error("CryptDestroyHash");
        return false;
    }
    hHash= 0;

    ByteVector buffer; buffer.resize(0x40000);

    DWORD crcsum= 0;
    DWORD filecrc= 0;

    bool bFlag= true;
    do {
        DWORD dwDataLen= fread(vectorptr(buffer), 1, bFlag?0x40:0x40000, fin);

        if (!CryptDecrypt(hKey, NULL, feof(fin), NULL, vectorptr(buffer), &dwDataLen))
        {
            error("CryptDecrypt");
            return false;
        }
        if (bFlag) {
            if (bVerbose) dump_nbf_header((char*)vectorptr(buffer));

            if (!(buffer[10]=='-' && buffer[31]=='-' && buffer[42]=='-' && buffer[52]=='-' && buffer[61]=='-'))
                printf("WARNING: this does not look like a nbf header, possible you provided the wrong password\n");
            else
                filecrc=strtoul((char*)vectorptr(buffer)+0x35,0,16);
        }

        crcsum += calccrc32(vectorptr(buffer), bFlag?0x34:dwDataLen);

        bFlag= false;
        fwrite(vectorptr(buffer), 1, dwDataLen, fout);
    } while (!feof(fin));

    fclose(fin);
    fclose(fout);
    CryptDestroyKey(hKey);
    CryptReleaseContext(hProv, 0);

    if (filecrc != crcsum)
        printf("WARNING: checksum is not ok, possibly you provided the wrong password\n");

    return true;
}
bool XorBuffer(BYTE *buf, int len, DWORD dwVal)
{
    DWORD *p= (DWORD*)buf;
    if (len&3)
        printf("WARNING: buffer not multiple of 4\n");
    len >>= 2;

    while (len--)
        *p++ ^= dwVal;

    return true;
}
bool xor_file(const char *szEncryptedFilename, const char *szDecryptedFilename, DWORD dwXorValue)
{
    FILE *fin= fopen(szEncryptedFilename, "rb");
    if (fin==NULL) {
        perror(szEncryptedFilename);
        return false;
    }
    FILE *fout= fopen(szDecryptedFilename, "wb");
    if (fout==NULL) {
        perror(szDecryptedFilename);
        return false;
    }
    ByteVector buffer; buffer.resize(0x40000);

    DWORD crcsum= 0;
    DWORD filecrc= 0;

    // both start out as false, when we have not yet determined what we are doing.
    bool fEncrypting= false;
    bool fDecrypting= false;

    bool bFlag= true;
    do {
        DWORD dwDataLen= fread(vectorptr(buffer), 1, bFlag?0x40:0x40000, fin);

        if (bFlag) {
            if (buffer[10]=='-' && buffer[31]=='-' && buffer[42]=='-' && buffer[52]=='-' && buffer[61]=='-') {
                fEncrypting= true;
                filecrc=strtoul((char*)vectorptr(buffer)+0x35,0,16);

                if (bVerbose) dump_nbf_header((char*)vectorptr(buffer));
            }
        }
        if (fEncrypting)
            crcsum += calccrc32(vectorptr(buffer), bFlag?0x34:dwDataLen);

        if (!XorBuffer(vectorptr(buffer), dwDataLen, dwXorValue))
        {
            error("XorBuffer");
            return false;
        }
        if (bFlag) {

            if (buffer[10]=='-' && buffer[31]=='-' && buffer[42]=='-' && buffer[52]=='-' && buffer[61]=='-') {
                fDecrypting= true;
                filecrc=strtoul((char*)vectorptr(buffer)+0x35,0,16);

                if (bVerbose) dump_nbf_header((char*)vectorptr(buffer));
            }

            if (fEncrypting==fDecrypting)
                printf("WARNING: this does not look like a nbf header, possible you provided the wrong password\n");
        }

        if (fDecrypting)
            crcsum += calccrc32(vectorptr(buffer), bFlag?0x34:dwDataLen);

        bFlag= false;
        fwrite(vectorptr(buffer), 1, dwDataLen, fout);
    } while (!feof(fin));

    fclose(fin);
    fclose(fout);

    if (filecrc != crcsum)
        printf("WARNING: checksum is not ok, possibly you provided the wrong password\n");

    return true;
}

bool crccalc_file(char *szFilename, bool bDoUpdateCrc)
{
    FILE *fin= fopen(szFilename, "r+b");
    if (fin==NULL) {
        perror(szFilename);
        return false;
    }
    ByteVector buffer; buffer.resize(0x40000);

    DWORD crcsum= 0;
    DWORD filecrc= 0;

    bool bFlag= true;
    do {
        DWORD dwDataLen= fread(vectorptr(buffer), 1, bFlag?0x40:0x40000, fin);
        if (bFlag) {
            if (bVerbose) dump_nbf_header((char*)vectorptr(buffer));

            filecrc=strtoul((char*)vectorptr(buffer)+0x35,0,16);
        }

        crcsum += calccrc32(vectorptr(buffer), bFlag?0x34:dwDataLen);

        bFlag= false;
    } while (!feof(fin));

    if (bDoUpdateCrc) {
        char hexcrc[9];
        sprintf(hexcrc, "%08lx", crcsum);

        fseek(fin, 0x35, SEEK_SET);
        fwrite(hexcrc, 1, 8, fin);
        fseek(fin, 0, SEEK_END);

        printf("updated checksum to %08lx\n", crcsum);
    }
    else if (filecrc != crcsum) {
        printf("WARNING: checksum is not ok, calculated crcsum = %08lx    file says: %08lx\n", crcsum, filecrc);
    }

    fclose(fin);

    return true;
}
bool print_file_info(char *szFilename)
{
    FILE *fin= fopen(szFilename, "r+b");
    if (fin==NULL) {
        perror(szFilename);
        return false;
    }
    ByteVector buffer; buffer.resize(0x40);

    DWORD dwDataLen= fread(vectorptr(buffer), 1, 0x40, fin);
    dump_nbf_header((char*)vectorptr(buffer));

    fclose(fin);

    return true;
}


void setdata(char *dst, int len, char* src)
{
    while (len--) {
        if (*src) {
            *dst++= *src++;
        }
        else {
            *dst++= 0;
        }
    }
}
bool update_header(char *szFilename, char* szDevicetype, char* szOperatorname, char* szLanguage, char* szVersion)
{
    FILE *fin= fopen(szFilename, "r+b");
    if (fin==NULL) {
        perror(szFilename);
        return false;
    }
    ByteVector buffer; buffer.resize(0x40000);

    DWORD crcsum= 0;
    DWORD filecrc= 0;

    bool bFlag= true;
    do {
        DWORD dwDataLen= fread(vectorptr(buffer), 1, bFlag?0x40:0x40000, fin);

        if (bFlag) {
            if (bVerbose) dump_nbf_header((char*)vectorptr(buffer));

            if (szDevicetype) setdata((char*)vectorptr(buffer), 10, szDevicetype);
            if (szOperatorname) setdata((char*)vectorptr(buffer)+11, 20, szOperatorname);
            if (szLanguage) setdata((char*)vectorptr(buffer)+32, 10, szLanguage);
            if (szVersion) setdata((char*)vectorptr(buffer)+43, 9, szVersion);

            fseek(fin, 0, SEEK_SET);
            fwrite(vectorptr(buffer), 1, 0x35, fin);
            fseek(fin, 0x40, SEEK_SET);
        }

        crcsum += calccrc32(vectorptr(buffer), bFlag?0x34:dwDataLen);

        bFlag= false;
    } while (!feof(fin));

    char hexcrc[9];
    sprintf(hexcrc, "%08lx", crcsum);

    fseek(fin, 0x35, SEEK_SET);
    fwrite(hexcrc, 1, 8, fin);
    fseek(fin, 0, SEEK_END);

    fclose(fin);

    return true;
}
bool FindXorKey(const std::string& filename, DWORD dwOffset, DWORD dwValue, DWORD dwRepeat, DWORD& dwXorKey)
{
    FILE *fin= fopen(filename.c_str(), "rb");
    if (fin==NULL) {
        perror(filename.c_str());
        return false;
    }

    fseek(fin, dwOffset, SEEK_SET);

    DwordVector dwBuf; dwBuf.resize(dwRepeat);
    DWORD dwLen= fread(vectorptr(dwBuf), sizeof(DWORD), dwBuf.size(), fin);
    fclose(fin);

    if (dwLen!=dwBuf.size())
        return false;

    // all 4 values should be the same.
    for (int i= 1 ; i<dwBuf.size() ; i++) {
        if (dwBuf[i]!=dwBuf[0])
            return false;
    }

    dwXorKey= dwBuf[0]^dwValue;
    return true;
}
bool PrintXorKeyForFile(const std::string& filename, DWORD dwOffset, DWORD dwValue, DWORD dwRepeat, bool bDoDecrypt)
{
    DWORD dwXorKey;
    if (!FindXorKey(filename, dwOffset, dwValue, dwRepeat, dwXorKey))
        return false;

    std::string basename= filename;
    basename.resize(basename.find('.'));
    basename += ".nba";
    printf("xda2nbftool -x %-10s %-10s 0x%08lx\n", filename.c_str(), basename.c_str(), dwXorKey);

    if (bDoDecrypt) {
        xor_file(filename.c_str(), basename.c_str(), dwXorKey);
    }

    return true;
}
void FindXorKeys(bool bDoDecrypt)
{
    // nk.nbf has several problems with auto find:
    // 40-40040 can contain one of the following:
    //    - a bootloader
    //        * 0x80, 0x43454345, 1
    //    - filled with zeroes
    //        * 0x40, 0, 4
    //    - ??? can it be filled with ones?
    // 40040  contains the os image
    //    - the only known header word is the ECEC part, which does not
    //      provide much verification.
    //        * 0x40080, 0x43454345, 1
    //
    // another possibility: 0x40000, 0, 4 - assuming the bl is always 0 padded.
    PrintXorKeyForFile("NK.nbf", 0x40000, 0, 4, bDoDecrypt);

    PrintXorKeyForFile("ms_.nbf", 0x80, 0xffffffff, 4, bDoDecrypt);
    PrintXorKeyForFile("Radio_.nbf", 0x60, 0xffffffff, 4, bDoDecrypt);
}

void usage()
{
    printf("Usage: xda2nbftool [-u] { -c | -e | -d }  inputfile outputfile password\n");
    printf("    -c  : calculate crcsum of file\n");
    printf("    -u  : updates the calculated checksum\n");
    printf("    -e  : encrypt inputfile\n");
    printf("    -d  : decrypt inputfile\n");
    printf("    -x  : xor inputfile with password ( should be a number )\n");
    printf("    -t  : automatically find the xor key, with -x auto decrypt\n");
    printf("    -sd : change devicetype\n");
    printf("    -so : change operatorname\n");
    printf("    -sl : change language\n");
    printf("    -sv : change version\n");
    printf("    -v  : verbose\n");
    printf("\n");
    printf("    HimaUpgradeUt v0400, v0500 use password HTC\n");
    printf("    HimaUpgradeUt v0600 uses password 921211\n");
    printf("    HimaUpgradeUt v0910 uses xor with 0x89124137, 0x25863614 or 0x12345678\n");
    printf("    HimaUpgradeUt v1000 uses xor with 0x20040304, 0x20040305 or 0x20040306\n");
}
#define HANDLESTR2(str) (argv[i][3]? str= argv[i]+3 : i+1<argc ? str= argv[++i] : 0)
#define HANDLEULOPTION1(var, type) (argv[i][2] ? var= (type)strtoul(argv[i]+2, 0, 0) : i+1<argc ? var= (type)strtoul(argv[++i], 0, 0) : 0)

int main(int argc, char **argv)
{
    bool bDoEncrypt= false;
    bool bDoDecrypt= false;
    bool bDoCalcCrc= false;
    bool bDoUpdateCrc= false;
    bool bDoUpdateHdr= false;
    bool bDoXorDecrypt= false;
    bool bDoFindDecryption= false;
    DWORD dwXorValue= 0;

    char *new_devicetype= NULL;
    char *new_operatorname= NULL;
    char *new_language= NULL;
    char *new_version= NULL;

    char *args[3];

    int argsfound=0; 
    for (int i=1 ; i<argc ; i++)
    {
        if (argv[i][0]=='-') switch (argv[i][1])
        {
            case 'e': bDoEncrypt= true; break;
            case 'd': bDoDecrypt= true; break;
            case 'x': bDoXorDecrypt= true; break;
            case 'c': bDoCalcCrc= true; break;
            case 'u': bDoUpdateCrc= true; break;
            case 's': switch(argv[i][2]) {
                case 'd': HANDLESTR2(new_devicetype);  bDoUpdateHdr= true; break;
                case 'o': HANDLESTR2(new_operatorname);  bDoUpdateHdr= true; break;
                case 'l': HANDLESTR2(new_language);  bDoUpdateHdr= true; break;
                case 'v': HANDLESTR2(new_version);  bDoUpdateHdr= true; break;
                default:
                          usage();
                          return 1;
            }
                break;
            case 't': bDoFindDecryption= true; break;
            case 'v': bVerbose= true; break;
            default:
                usage();
                return 1;
        }
        else  {
            if (argsfound<3) {
                args[argsfound]= argv[i];
            }
            else {
                printf("too many args\n");
                usage();
                return 1;
            }
            argsfound++;
        }
    }
    if (!bDoFindDecryption && argsfound==0) {
        usage();
        return 0;
    }

    if (bDoFindDecryption) {
        FindXorKeys(bDoXorDecrypt);
    }
    else if (bDoXorDecrypt) {
        if (argsfound!=3) { usage(); return 1; }
        xor_file(args[0], args[1], strtoul(args[2], 0, 0));
    }
    else if (bDoDecrypt) {
        if (argsfound!=3) { usage(); return 1; }
        decrypt_file(args[0], args[1], args[2]);
    }
    else if (bDoEncrypt) {
        if (argsfound!=3) { usage(); return 1; }
        encrypt_file(args[0], args[1], args[2]);
    }
    else if (bDoCalcCrc && !(bDoEncrypt||bDoDecrypt)) {
        if (argsfound!=1) { usage(); return 1; }
        crccalc_file(args[0], bDoUpdateCrc);
    }
    else if (bDoUpdateHdr) {
        if (argsfound!=1) { usage(); return 1; }
        update_header(args[0], new_devicetype, new_operatorname, new_language, new_version);
    }
    else if (bVerbose) {
        if (argsfound!=1) { usage(); return 1; }
        print_file_info(args[0]);
    }
}