Python3 实现Base64

• 支持自定义64字符

def b64_encode(data, *, letters=None):
    import struct
    if not isinstance(data, bytes):
        raise ValueError('type of the data which to be encoded must be bytes')
    if letters is None:
        letters = b'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'
    if not isinstance(letters, bytes):
        raise ValueError('type of the letters which to be trans must be bytes')
    if len(letters) != 64:
        raise ValueError('letters length must 64')

    res = bytearray()

    for i in range(0, len(data), 3):
        v = data[i:i + 3]
        v_len = len(v)
        if v_len == 3:
            n = struct.unpack('>I', b'\x00' + v)[0]
            res.append(letters[n >> 18 & 0x3f])
            res.append(letters[n >> 12 & 0x3f])
            res.append(letters[n >> 6 & 0x3f])
            res.append(letters[n >> 0 & 0x3f])
        elif v_len == 1:
            n = struct.unpack('>I', b'\x00\x00' + v + b'\x00')[0]
            res.append(letters[n >> 10 & 0x3f])
            res.append(letters[n >> 4 & 0x3f])
        else:
            n = struct.unpack('>I', b'\x00' + v + b'\x00')[0]
            res.append(letters[n >> 18 & 0x3f])
            res.append(letters[n >> 12 & 0x3f])
            res.append(letters[n >> 6 & 0x3f])

    if len(res) % 4:
        res.extend(b'=' * (4 - len(res) % 4))
    return bytes(res)


def b64_decode(data, *, letters=None):
    import struct
    if not isinstance(data, bytes):
        raise ValueError('type of the data which to be encoded must be bytes')
    if letters is None:
        letters = b'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'
    if not isinstance(letters, bytes):
        raise ValueError('type of the letters which to be trans must be bytes')
    if len(letters) != 64:
        raise ValueError('letters length must 64')

    res = bytearray()
    data = data.rstrip(b'=')
    for i in range(0, len(data), 4):
        v = data[i:i + 4]
        v_len = len(v)
        if v_len == 4:
            n = 0
            for _k in range(4):
                n <<= 6
                n |= letters.index(v[_k])
            res.extend(n.to_bytes(3, byteorder='big'))
        elif v_len == 2:
            n = 0
            n |= letters.index(v[0])
            n <<= 2
            n |= (letters.index(v[1]) >> 4)
            res.extend(n.to_bytes(1, byteorder='big'))
        else:
        	# v_len == 3
            n = 0
            n |= letters.index(v[0])
            n <<= 6
            n |= letters.index(v[1])
            n <<= 4
            n |= (letters.index(v[2]) >> 2)
            res.extend(n.to_bytes(2, byteorder='big'))

    return bytes(res)


def test(f1, f2, d1, d2):
    import os
    for i in range(500):
        for j in range(50):
            data = os.urandom(i)
            v1 = f1(data)
            v2 = f2(data)
            if not v1 == v2:
                print(data, v1, v2)
                return
            dd1 = d1(v2)
            dd2 = d2(v1)
            if not (dd1 == dd2 == data):
                print(data, v1, v2, dd1, dd2)
                return


def test_letters():
    import random
    import os
    letters = b'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'
    b = bytearray(letters)
    random.shuffle(b)
    b = bytes(b)
    print('letters', b)

    for i in range(500):
        for j in range(50):
            data = os.urandom(i)
            if not (b64_decode(b64_encode(data)) == data):
                print('Error', data)
                return


if __name__ == '__main__':
    import base64
    test(b64_encode, base64.b64encode, b64_decode, base64.b64decode)
    test_letters()