xof_stream_cipher_b64.py

"""XOF Stream Cipher: encrypt and decrypt strings with shake256 - base64 version"""


from base64 import b64encode, b64decode
from hashlib import new
from secrets import token_bytes
from sys import argv


def xor(b1: bytes, b2: bytes) -> bytes:
    """XOR two equal-length byte strings together."""
    b3 = bytearray()
    for i in range(len(b1)):
        b3.append(b1[i] ^ b2[i])

    return bytes(b3)


def encrypt(key: bytes, plaintext: bytes) -> bytes:
    """Stream cipher encrypt."""
    key_bytes = new('shake256', key).digest(len(plaintext))

    return xor(plaintext, key_bytes)


def decrypt(key: bytes, ciphertext: bytes) -> bytes:
    """Stream cipher decrypt."""
    key_bytes = new('shake256', key).digest(len(ciphertext))

    return xor(ciphertext, key_bytes)


def hmac(key: bytes, message: bytes) -> bytes:
    """Create an hmac according to rfc 2104 specifications."""
    # set up variables
    B, L = 136 , len(message)
    L = L if L < 32 else 32
    ipad_byte = 0x36.to_bytes(1, 'big')
    opad_byte = 0x5c.to_bytes(1, 'big')
    null_byte = 0x00.to_bytes(1, 'big')
    ipad = b''.join([ipad_byte for i in range(B)])
    opad = b''.join([opad_byte for i in range(B)])

    # if key length is greater than digest length, hash it first
    key = key if len(key) <= L else new('shake256', key).digest(L)

    # if key length is less than block length, pad it with null bytes
    key = key + b''.join(null_byte for i in range(B - len(key)))

    # compute and return the hmac
    partial = new('shake256', xor(key, ipad) + message).digest(L)
    return new('shake256', xor(key, opad) + partial).digest(L)


def check_hmac(key: bytes, message: bytes, mac: bytes) -> bool:
    """Check an hmac."""
    # first compute the proper hmac
    computed = hmac(key, message)

    # if it is the wrong length, reject
    if len(mac) != len(computed):
        return False

    # compute difference without revealing anything through timing attack
    diff = 0
    for i in range(len(mac)):
        diff += mac[i] ^ computed[i]

    return diff == 0


def seal(key: bytes, plaintext: bytes, iv_size: int = 32) -> str:
    """Generate an iv, encrypt a message, and create an hmac all in one."""
    iv = token_bytes(iv_size) if iv_size < len(plaintext) else token_bytes(len(plaintext))
    ct = encrypt(key + iv, plaintext)
    return (b64encode(iv) + b'.' + b64encode(ct) + b'.' + b64encode(hmac(key, ct))).decode()


def unseal(key: bytes, ciphertext: str) -> bytes:
    """Checks hmac, then decrypts the message."""
    ciphertext = ciphertext.split('.')
    iv = b64decode(ciphertext[0])
    ct = b64decode(ciphertext[1])
    ac = b64decode(ciphertext[2])

    if not check_hmac(key, ct, ac):
        raise Exception('HMAC authentication failed')

    return decrypt(key + iv, ct)


def main(args):
    """Main function for invoking as cli tool."""
    # parse arguments
    mode = args[0]
    key = bytes(args[1], 'utf-8')
    text = args[2]

    if mode == 'encrypt' or mode == 'e':
        print(b64encode(encrypt(key, bytes(text, 'utf-8'))).decode())
    elif mode == 'decrypt' or mode == 'd':
        print(decrypt(key, b64decode(text)).decode())
    elif mode == 'seal' or mode == 's':
        print(seal(key, bytes(text, 'utf-8')))
    elif mode == 'open' or mode == 'o':
        print(unseal(key, text).decode())
    else:
        print('unknown mode: ' + mode)


if __name__ == '__main__':
    if len(argv) < 4:
        print('usage: ' + argv[0] + ' [encrypt|e|decrypt|d|seal|s|open|o] [key] [plaintext|ciphertext]')
        exit(1)
    main(argv[1:])


def license():
    """Copyleft (c) 2021 k98kurz

        Permission to use, copy, modify, and/or distribute this software
        for any purpose with or without fee is hereby granted, provided
        that the above copyleft notice and this permission notice appear in
        all copies.

        THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
        WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
        WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
        AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR
        CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
        OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
        NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
        CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
    """
    return license.__doc__