sqlite3采用MD5和Aes联合加密技术

来源:互联网 发布:多益网络的游戏客服 编辑:程序博客网 时间:2024/05/16 11:34
#include "aes/aes.h"
#include "md5/md5.h"
#include "src/sqlite3.c"


#define SQLITECRYPTERROR_PROVIDER "Cryptographic provider not available"

typedef unsigned char byte_t;


struct CryptBlock
{
    void*        pBuffer;
    int            bufSize;
    AES_KEY        key;
    int            need_key;
};

const byte_t g_ivec[] = {
    0x05,0x66,0xfa,0xa4,
    0x32,0x4f,0x1f,0xa2,
    0x34,0x67,0x33,0x9b,
    0xb6,0x7e,0x3d,0xc2
};

const byte_t g_ecount[] = {
    0xa1,0xc6,0x4a,0x74,
    0x3d,0xc7,0x8e,0x54,
    0x91,0xce,0xe3,0xbc,
    0xc2,0xa2,0xd3,0xb9
};


static void * sqlite3pager_get_codecarg(Pager *pPager)
{
  return (pPager->xCodec) ? pPager->pCodec : NULL;
}

static int create_aes_key(char* pKey,int nKey,AES_KEY* paKey)
{
    byte_t aes_key[MD5_DIGEST_LENGTH];
    MD5_CTX ctx;

    memset(aes_key,0,MD5_DIGEST_LENGTH);
    MD5_Init(&ctx);
    MD5_Update(&ctx,pKey,nKey);
    MD5_Final(aes_key,&ctx);


    return AES_set_encrypt_key(aes_key,KEY_128,paKey);

}

static void encrypt(void* data,struct CryptBlock* pBlock)
{
    size_t nPageSize = 0;
    byte_t iv[AES_BLOCK_SIZE];
    byte_t ecount_buf[AES_BLOCK_SIZE];
    int num = 0;

    if (!pBlock->need_key) return;

    memcpy(iv,g_ivec,AES_BLOCK_SIZE);
    memcpy(ecount_buf,g_ecount,AES_BLOCK_SIZE);
        
    nPageSize = pBlock->bufSize;
    AES_ctr128_encrypt((byte_t*)data, (byte_t*)pBlock->pBuffer,
        nPageSize, &pBlock->key,iv,ecount_buf,&num);
}

#define DATA_TO_PGHDR(D)  (&((PgHdr*)(D))[-1])

// Encrypt/Decrypt functionality, called by pager.c
void * sqlite3Codec(void *pArg, void *data,Pgno nPageNum, int nMode)
{
     struct CryptBlock* pBlock = (struct CryptBlock*)pArg;


     if( !pBlock ) return data;

     if(nMode != 2)
     {
          
         DbPage* pDbPage = 0;
         pDbPage = DATA_TO_PGHDR(data);
        
         if(pDbPage->pPager->pageSize != pBlock->bufSize)
         {
             sqlite3_free(pBlock->pBuffer);
             pBlock->pBuffer = sqlite3_malloc(pDbPage->pPager->pageSize);
             pBlock->bufSize = pDbPage->pPager->pageSize;
         }
     }

     switch(nMode)
     {
     case 0: // Undo a "case 7" journal file encryption
     case 2: // Reload a page
     case 3: // Load a page
        encrypt(data,pBlock);
        memcpy(data,pBlock->pBuffer,pBlock->bufSize);
        return data;
     case 6: // Encrypt a page for the main database file
     case 7:
        encrypt(data,pBlock);
        break;
     }

     return pBlock->pBuffer;
}

void sqlite3CodecFree(void* pCodecArg)
{
    struct CryptBlock* pBlock = (struct CryptBlock*)pCodecArg;
    if(pBlock)
    {
        if(pBlock->pBuffer)
        {
            sqlite3_free(pBlock->pBuffer);
            pBlock->pBuffer = 0;
        }
    }
}

void sqlite3CodecSizeChng(void* pCodecArg,int size,int reserve)
{
    struct CryptBlock* pBlock = (struct CryptBlock*)pCodecArg;

    if(pBlock && pBlock->pBuffer && pBlock->bufSize != size)
    {
        sqlite3_free(pBlock->pBuffer);
        pBlock->pBuffer = sqlite3_malloc(size);
        pBlock->bufSize = size;
    }
}


int sqlite3_key(
  sqlite3 *db,                   /* Database to be rekeyed */
  const void *pKey, int nKey     /* The key */
)
{
    return sqlite3CodecAttach(db, 0, pKey, nKey);
}

int sqlite3_rekey(
  sqlite3 *db,                   /* Database to be rekeyed */
  const void *pKey, int nKey     /* The new key */
)
{
    Btree *pbt = db->aDb[0].pBt;
    Pager *p = sqlite3BtreePager(pbt);
    struct CryptBlock* pBlock = (struct CryptBlock*)sqlite3pager_get_codecarg(p);
    AES_KEY* aKey = 0;
    int rc = SQLITE_ERROR;


    if (!pBlock && !pBlock->need_key) return SQLITE_OK; // Wasn't encrypted to begin with

    aKey = &pBlock->key;


    sqlite3PagerSetCodec(p,sqlite3Codec,
            sqlite3CodecSizeChng,sqlite3CodecFree,
            pBlock);

    // Start a transaction
    rc = sqlite3BtreeBeginTrans(pbt, 1);

    if (!rc)
    {
        // Rewrite all the pages in the database using the new encryption key
        Pgno nPage = 0;
        Pgno nSkip = PAGER_MJ_PGNO(p);
        DbPage *pPage;
        Pgno n;

        sqlite3PagerPagecount(p,&nPage);
        for(n = 1; rc == SQLITE_OK && n <= nPage; n ++)
        {
            if (n == nSkip) continue;
            rc = sqlite3PagerGet(p, n, &pPage);
            if(!rc)
            {
                rc = sqlite3PagerWrite(pPage);
                sqlite3PagerUnref(pPage);
            }
        }
    }

    // If we succeeded, try and commit the transaction
    if (!rc)
    {    
        rc = sqlite3BtreeCommit(pbt);
    }

    // If we failed, rollback
    if (rc)
    {
        sqlite3BtreeRollback(pbt);
    }

    sqlite3PagerSetCodec(p,0,0,0,0);
    sqlite3_free(pBlock);
    
  return rc;
}

void sqlite3_activate_see(
  const char *zPassPhrase        /* Activation phrase */
)
{
    
}

void sqlite3CodecGetKey(sqlite3 *db, int nDb, void **ppKey, int *pnKeyLen)
{
  *ppKey = 0;
  *pnKeyLen = 0;
}

int sqlite3CodecAttach(sqlite3 *db, int nDb, const void *pKey, int nKeyLen)
{
    int rc = SQLITE_ERROR;
    AES_KEY* a_key = 0;
    struct CryptBlock* pBlock = 0;

    // No key specified, could mean either use the main db's encryption or no encryption
    if (!pKey || !nKeyLen)
    {
        if (!nDb)
        {
          return SQLITE_OK; // Main database, no key specified so not encrypted
        }
        else
        {
            //获取主数据库的加密块并复制密钥给附加数据库使用
            pBlock = (struct CryptBlock*)sqlite3pager_get_codecarg(sqlite3BtreePager(db->aDb[0].pBt));
            if(!pBlock) return SQLITE_OK;
            if(!pBlock->need_key) return SQLITE_OK;

            a_key = &pBlock->key;
        }
    }
    else
    {
        pBlock = sqlite3_malloc(sizeof(struct CryptBlock));
        a_key = &pBlock->key;
        pBlock->bufSize = 0;
        pBlock->need_key = 1;
        pBlock->pBuffer = 0;
        if(create_aes_key((char*)pKey,nKeyLen,a_key))
        {
            sqlite3Error(db, rc, SQLITECRYPTERROR_PROVIDER);
            return rc;
        }
    }

    if(a_key)
    {
        Pager* pPager = sqlite3BtreePager(db->aDb[nDb].pBt);
        
        sqlite3PagerSetCodec(pPager,sqlite3Codec,
            sqlite3CodecSizeChng,0,
            pBlock);
        
    }

    return SQLITE_OK;
}