box-o-sand/onetimepad.py

from itertools import izip

PAD_MODULO = 3
PADLINE_OFFSET = -2
TEXTLINE_OFFSET = -1
DEFAULT_PAD_WIDTH = 72
CIPHER_FLOOR = 1
CIPHER_CEIL = 25
NULLCHAR = '.'


def encode(msg, pad):
    ret = []
    filled = list(_text_padfill(msg, pad))
    for i, line in enumerate(_text_cipherfill(filled)):
        lineno = i + 1
        if _is_cipherline(lineno):
            ret.append(line)
    return ''.join(ret).strip(NULLCHAR)


def decode(ciphertext, pad):
    ret = []
    filled = list(_cipher_padfill(ciphertext, pad))
    for i, line in enumerate(_cipher_textfill(filled)):
        lineno = i + 1
        if _is_textline(lineno):
            ret.append(line)
    return ''.join(ret).strip(NULLCHAR)


def create_pad(length, width=DEFAULT_PAD_WIDTH):
    return '\n'.join(create_pad_lines(length, width=width))


def create_pad_lines(length, width=DEFAULT_PAD_WIDTH):
    chars = _create_chars_for_pad(length)
    lines = _chunk_chars_into_lines(chars, width=width)
    for line in lines:
        yield line
        yield ''
        yield ''


def _create_chars_for_pad(length):
    chars = []
    for char in range(length):
        chars.append(_AS_ALPHA[_get_randint()])
    return ''.join(chars)


def _chunk_chars_into_lines(chars, width=DEFAULT_PAD_WIDTH):
    lines = []
    for chunk in _as_line_chunks(chars, width=width):
        lines.append(chunk)
    return lines


def _as_line_chunks(chars, width=DEFAULT_PAD_WIDTH):
    chunk = []
    for char in chars.replace('\n', ''):
        chunk.append(char)
        if len(chunk) == width:
            yield ''.join(chunk)
            chunk = []
    if len(chunk) < width:
        chunk += ([NULLCHAR] * (width - len(chunk)))
    yield ''.join(chunk)


def _text_padfill(text, padlines):
    if isinstance(padlines, basestring):
        padlines = padlines.splitlines()
    padlines = [line.strip() for line in padlines if line.strip()]
    textlines = _chunk_chars_into_lines(text, _get_textwidth(padlines))
    for textline, padline in izip(textlines, padlines):
        yield padline
        yield textline
        yield ''


def _cipher_padfill(ciphertext, padlines):
    if isinstance(padlines, basestring):
        padlines = padlines.splitlines()
    padlines = [line.strip() for line in padlines if line.strip()]
    cipherlines = _chunk_chars_into_lines(ciphertext, _get_textwidth(padlines))
    for cipherline, padline in izip(cipherlines, padlines):
        yield padline
        yield ''
        yield cipherline


def _text_cipherfill(padfilled_text_lines):
    for i, line in enumerate(padfilled_text_lines):
        lineno = i + 1
        if _is_cipherline(lineno):
            padline = padfilled_text_lines[i - abs(PADLINE_OFFSET)]
            textline = padfilled_text_lines[i - abs(TEXTLINE_OFFSET)]
            yield padline
            yield textline
            yield _cipherline_from_padline_and_textline(padline, textline)


def _cipher_textfill(padfilled_cipher_lines):
    for i, line in enumerate(padfilled_cipher_lines):
        lineno = i + 1
        if _is_cipherline(lineno):
            padline = padfilled_cipher_lines[i - abs(PADLINE_OFFSET)]
            textline = padfilled_cipher_lines[i - abs(TEXTLINE_OFFSET)]
            yield padline
            yield _textline_from_cipherline_and_padline(line, padline)
            yield line


def _cipherline_from_padline_and_textline(padline, textline):
    ret = []
    for padchar, textchar in izip(padline, textline):
        if textchar == NULLCHAR:
            ret.append(NULLCHAR)
            continue
        charnum = _AS_NUMS[padchar] + _AS_NUMS[textchar]
        idx = charnum if charnum <= CIPHER_CEIL else charnum % CIPHER_CEIL
        ret.append(_AS_ALPHA[idx])
    return ''.join(ret)


def _textline_from_cipherline_and_padline(cipherline, padline):
    ret = []
    for ciphercar, padchar in izip(cipherline, padline):
        if ciphercar == NULLCHAR:
            ret.append(NULLCHAR)
            continue
        charnum = _AS_NUMS[ciphercar] - _AS_NUMS[padchar]
        idx = charnum if charnum <= CIPHER_CEIL else charnum % CIPHER_CEIL
        if idx < 0:
            idx = CIPHER_CEIL + idx
        ret.append(_AS_ALPHA[idx])
    return ''.join(ret)


def _get_textwidth(textlines):
    return max([len(line) for line in textlines])


def _is_padline(lineno):
    return _is_cipherline(lineno + abs(PADLINE_OFFSET))


def _is_textline(lineno):
    return _is_cipherline(lineno + abs(TEXTLINE_OFFSET))


def _is_cipherline(lineno):
    return not lineno % PAD_MODULO


def _get_randint(randomness_file='/dev/urandom', fallback_to_fake=True):
    if not _DEV_RANDOM.get('fp'):
        _get_urandom_fp(randomness_file=randomness_file,
                        fallback_to_fake=fallback_to_fake)
    ret = ord(_DEV_RANDOM['fp'].read(1)) % CIPHER_CEIL
    return ret if ret != 0 else _get_randint()


class _FakeDevUrandom(object):

    def read(self, nchars):
        import random
        return ''.join([chr(random.randint(0, 255)) for i in range(0, nchars)])


def _get_urandom_fp(randomness_file='/dev/urandom', fallback_to_fake=True):
    try:
        _DEV_RANDOM['fp'] = open(randomness_file)
    except (OSError, IOError):
        if fallback_to_fake:
            _DEV_RANDOM['fp'] = _FakeDevUrandom()
        else:
            raise


_DEV_RANDOM = {}


def _mk_as_alpha_as_nums():
    as_alpha = {}
    as_nums = {}
    a_chr = ord('A')
    past_j = False

    for char in range(CIPHER_CEIL + 1):
        letter = chr(a_chr + char)
        if letter != 'J':
            key = char + (1 if not past_j else 0)
            as_alpha[key] = letter
        else:
            past_j = True

    for key, val in as_alpha.iteritems():
        as_nums[val] = key

    return as_alpha, as_nums


_AS_ALPHA, _AS_NUMS = _mk_as_alpha_as_nums()