// Copyright: (c) Magnus Therning, 2012
// License: BSD3, found in the LICENSE file

#include <assert.h>

#include "b64.h"

static char const encmap[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

void b64_enc_part(uint8_t const *src, size_t srclen,
    uint8_t *dst, size_t *dstlen,
    uint8_t const **rem, size_t *remlen)
{
    assert(src || 0 == srclen);
    assert(dst);
    assert(dstlen);
    assert(rem);
    assert(remlen);

    size_t od = *dstlen, i;

    for(i = 0, *dstlen = 0; i + 3 <= srclen && *dstlen + 4 <= od; i += 3, *dstlen += 4) {
        int32_t o0, o1, o2, o3;
        o0 = src[i] >> 2;
        o1 = ((src[i] << 4) | (src[i+1] >> 4)) & 0x3f;
        o2 = ((src[i+1] << 2) | (src[i+2] >> 6)) & 0x3f;
        o3 = src[i+2] & 0x3f;
        *dst++ = encmap[o0];
        *dst++ = encmap[o1];
        *dst++ = encmap[o2];
        *dst++ = encmap[o3];
    }

    *rem = src + i;
    *remlen = srclen - i;
}

int b64_enc_final(uint8_t const *src, size_t srclen,
    uint8_t *dst, size_t *dstlen)
{
    assert(src || 0 == srclen);
    assert(dst);
    assert(dstlen);

    switch(srclen) {
        int32_t o0, o1, o2, o3;
    case 0:
        *dstlen = 0;
        return(0);
        break;
    case 1:
        o0 = src[0] >> 2;
        o1 = (src[0] << 4) & 0x3f;
        *dst++ = encmap[o0];
        *dst++ = encmap[o1];
        *dst++ = '=';
        *dst++ = '=';
        *dstlen = 4;
        return(0);
        break;
    case 2:
        o0 = src[0] >> 2;
        o1 = ((src[0] << 4) | (src[1] >> 4)) & 0x3f;
        o2 = (src[1] << 2) & 0x3f;
        *dst++ = encmap[o0];
        *dst++ = encmap[o1];
        *dst++ = encmap[o2];
        *dst++ = '=';
        *dstlen = 4;
        return(0);
        break;
    default:
        return(1);
        break;
    }
}

// decode map, 0x80 = not allowed, 0x40 = end char
static uint8_t const decmap[] = {
    0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
    0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
    0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x3e, 0x80, 0x80, 0x80, 0x3f,
    0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3a, 0x3b, 0x3c, 0x3d, 0x80, 0x80, 0x80, 0x40, 0x80, 0x80,
    0x80, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
    0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x80, 0x80, 0x80, 0x80, 0x80,
    0x80, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, 0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28,
    0x29, 0x2a, 0x2b, 0x2c, 0x2d, 0x2e, 0x2f, 0x30, 0x31, 0x32, 0x33, 0x80, 0x80, 0x80, 0x80, 0x80,
    0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
    0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
    0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
    0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
    0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
    0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
    0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
    0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80 };


int b64_dec_part(uint8_t const *src, size_t srclen,
    uint8_t *dst, size_t *dstlen,
    uint8_t const **rem, size_t *remlen)
{
    assert(src || 0 == srclen);
    assert(dst);
    assert(dstlen);
    assert(rem);
    assert(remlen);

    size_t od = *dstlen, i;
    int res = 0;

    for(i = 0, *dstlen = 0; i + 4 <= srclen && *dstlen + 3 <= od; i += 4, *dstlen += 3) {
        uint8_t o0, o1, o2, o3;

        o0 = decmap[src[i]];
        o1 = decmap[src[i+1]];
        o2 = decmap[src[i+2]];
        o3 = decmap[src[i+3]];
        if(!(0xc0 & (o0 | o1 | o2 | o3))) { // no illegal chars, and no '='
            *dst++ = (o0 << 2) | (o1 >> 4);
            *dst++ = (o1 << 4) | (o2 >> 2);
            *dst++ = (o2 << 6) | o3;
        } else if(!(0xc0 & (o0 | o1)) && (0x40 & o2 & o3) // two legal chars, two '='
            || !(0xc0 & (o0 | o1 | o2)) && (0x40 & o3)) { // three legal chars, one '='
            res = 0;
            break;
        } else {
            res = 1;
            break;
        }
    }

    *rem = src + i;
    *remlen = srclen - i;
    return(res);
}

int b64_dec_final(uint8_t const *src, size_t srclen,
    uint8_t *dst, size_t *dstlen)
{
    assert(src || 0 == srclen);
    assert(dst);
    assert(dstlen);

    uint8_t o0, o1, o2, o3;

    if(0 == srclen) {
        *dstlen = 0;
        return(0);
    }
    o0 = decmap[src[0]];
    o1 = decmap[src[1]];
    o2 = decmap[src[2]];
    o3 = decmap[src[3]];
    if(!(0xc0 & (o0 | o1)) && (0x40 & o2 & o3)) { // two legal chars, two '='
        *dst++ = (o0 << 2) | (o1 >> 4);
        *dstlen = 1;
    } else if(!(0xc0 & (o0 | o1 | o2)) && (0x40 & o3)) { // three legal chars, one '='
            *dst++ = (o0 << 2) | (o1 >> 4);
            *dst++ = (o1 << 4) | (o2 >> 2);
            *dstlen = 2;
    } else
        return(1);

    return(0);
}