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

// tool to write lotsof or very large files, with recognizable block signatures
//    - useful for determining how a flash drive erases files / diskspace
class rngsrc_mscryptoapi  {
    HCRYPTPROV hProv;
public:
    rngsrc_mscryptoapi()
    {
        if (!CryptAcquireContext(&hProv, NULL, NULL, PROV_RSA_FULL, CRYPT_VERIFYCONTEXT))
            throw "CryptAcquireContext error";
    }
    ~rngsrc_mscryptoapi()
    {
        CryptReleaseContext(hProv, 0);
    }
    void getentropy(uint8_t*p, size_t n)
    {

        if (!CryptGenRandom(hProv, n, (BYTE*)p))
            throw "CryptGenRandom error";

    }
};
void put32le(uint8_t*p, uint32_t x)
{
    p[0]= x;
    p[1]= x>>8;
    p[2]= x>>16;
    p[3]= x>>24;
}
void makebuffer(int gen, int blk, uint8_t *p, size_t n)
{
    static rngsrc_mscryptoapi  rng;
    DWORD ts= GetTickCount();
    SYSTEMTIME now;
    GetLocalTime(&now);
    FILETIME ft;
    if (!SystemTimeToFileTime(&now, &ft))
        error("setnow:SystemTimeToFileTime");

    // block id
    put32le(p+0, gen);
    put32le(p+4, blk);
    // timestamp
    put32le(p+8, ts);
    put32le(p+12, ft.dwLowDateTime);
    put32le(p+16, ft.dwHighDateTime);
    // random constant
    put32le(p+20, 0x76334872);
    put32le(p+24, 0x61626d62);
    rng.getentropy(p+28, n-28);
}

bool writelarge(int gen)
{
    FFHANDLE h;
    if (!OpenFileForWriting(std::string("large.bin"), h))
        return false;

    int blk=0;
    ByteVector buf(512);
    while (true)
    {
        makebuffer(gen, blk, &buf[0], buf.size());
        DWORD wrote;
        if (!WriteFile(h, &buf[0], buf.size(), &wrote, NULL))
            break;

        blk++;
    }
    CloseFile(h);
}
bool writesmall(int gen, int blk)
{
    FFHANDLE h;
    if (!OpenFileForWriting(std::string("small.bin"), h))
        return false;

    ByteVector buf(512);

    makebuffer(gen, blk, &buf[0], buf.size());
    DWORD wrote;
    if (!WriteFile(h, &buf[0], buf.size(), &wrote, NULL))
        ;

    CloseFile(h);
}
void usage()
{
    debug("Usage: writelarge [-b BLKSIZE] [-l] [-s]\n");
}
int WINAPI WinMain( HINSTANCE hInstance,
                   HINSTANCE hPrevInstance,
                   LPTSTR    lpCmdLine,
                   int       nCmdShow)
{
    DebugSetLogfile("writelarge.log");
    StringList args;
    if (!SplitString(ToString(lpCmdLine), args, false))
    {
        error("Error in commandline");
        return false;
    }

    size_t blksperfle= 1;   // -1 == until disk is full.
    bool createdirs= false; // create subdirs x/x/x/x/file
    bool cleanafter= false; // clean files after each generation

    size_t blksize;
    for (StringList::iterator i=args.begin() ; i!=args.end() ; ++i)
    {
        if (i->size()>1 && i->at(0)=='-')
            switch(i->at(1))
            {
                case 'b': getarg(i, args.end(), blksize); break;
                case 'n': getarg(i, args.end(), blksperfle); break;
                case 'd': createdirs=true; break;
                case 'c': cleanafter=true; break;
            }
        else
            usage();
    }


    if (args.front()=="small") {
        for (int gen=0 ; gen<4 ; gen++)
        {
            for (int n=0 ; n<256*1024 ; n++)
                writesmall(gen, n);
        }
    }
    else if (args.front()=="large") {
        for (int gen=0x10000 ; gen<0x10004 ; gen++)
        {
            writelarge(gen);
        }
    }
    else {
        debug("Usage: writelarge <large|small>\n");
    }
    return 0;
}