Bitwise XOR an input stream with a known key

xor.py

#!/usr/bin/env python

import argparse
import sys

from builtins import bytes

__title__ = 'xor.py'
__version__ = '1.0'
__author__ = 'malvidin'
__description__ = 'Bitwise XOR a file input stream with a known key'


def auto_int(x):
    return int(x, 0)


def xor(inputstream, xorkey):
    inputstream = bytearray(inputstream)
    xorkey = bytearray(xorkey)
    # Both input and xor key should be bytes
    assert len(inputstream) > 0, 'Input cannot be empty, did user provide offsets that create a zero length input?'
    assert len(xorkey) > 0, 'Key cannot be empty'
    result = b''
    for i in range(len(inputstream)):
        j = i % len(xorkey)
        result += bytes([inputstream[i] ^ xorkey[j]])
    return result


def main():
    parser = argparse.ArgumentParser(description=__description__)
    parser.add_argument('-k', '--keytype', default='hex', choices=['str', 'hex'],
                        help='How the key is entered. String (spam) or hex (7370616d).')
    parser.add_argument('-o', '--output', default='-', help='Output file to write to, default to STDOUT (-)')
    parser.add_argument('key', help='Key to use to decrypt the file stream.')
    parser.add_argument('-a', '--begin', type=auto_int, default=None, help='Offset to begin in the file stream.')
    group = parser.add_mutually_exclusive_group()
    group.add_argument('-l', '--length', type=auto_int, help='Bytes of data to retrieve from the file stream.')
    group.add_argument('-z', '--end', type=auto_int, default=None, help='Offset to end in the file stream.')

    # Data from STDIN
    if not sys.stdin.isatty():
        args = parser.parse_args()
        # Read all data from STDIN
        if sys.version_info[0] > 2:
            data = sys.stdin.buffer.read()
        else:
            data = sys.stdin.read()

    # Data from filename argument
    else:
        parser.add_argument('inputfile', help='Filename of the file to process. Ignored if reading from STDIN.')
        args = parser.parse_args()
        # Read all data from file
        data = open(args.inputfile, 'rb').read()

    # Slice data, and calculate position of end of portion if length is provided
    if args.length:
        assert args.length > 0, "Bytes of data must be a positive integer"
        args.end = (int(args.begin or 0) + args.length)
    portion = data[args.begin:args.end]

    # Create the representation of the key based on the keytype
    if args.keytype == 'str':
        key = args.key
    else:
        key = bytearray.fromhex(args.key)

    # XOR the portion
    xorslice = xor(portion, key)

    # Output to STDOUT
    if args.output == '-':
        if sys.version_info[0] > 2:
            sys.stdout.buffer.write(xorslice)
        else:
            sys.stdout.write(xorslice)
    # Or to file
    else:
        with open(args.output, 'wb') as f:
            f.write(xorslice)


if __name__ == '__main__':
    main()