box-o-sand/one-time-pad/onetimepad.py

from itertools import izip

DEFAULT_PAD_WIDTH = 72


def mod25encode(msg, pad):
    mod25 = Mod25(pad=pad)
    return mod25.encode(msg)


def mod25decode(msg, pad):
    mod25 = Mod25(pad=pad)
    return mod25.decode(msg)


class Mod25(object):
    _pad_modulo = 3
    _padline_offset = -2
    _textline_offset = -1
    _cipher_floor = 1
    _cipher_ceil = 25
    _nullchar = '.'

    def __init__(self, pad=''):
        self.pad = pad
        self._as_alpha = {}
        self._as_nums = {}
        self._mk_as_alpha_as_nums()

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

        for char in range(self._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

        self._as_alpha.update(as_alpha)
        self._as_nums.update(as_nums)

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

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

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

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

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

    def _ensure_pad_is_lines(self):
        if isinstance(self.pad, basestring):
            self.pad = self.pad.splitlines()

    def _text_padfill(self, text):
        self._ensure_pad_is_lines()
        padlines = [line.strip() for line in self.pad if line.strip()]
        textlines = _chunk_chars_into_lines(text, self._nullchar,
                                            _get_textwidth(padlines))
        for textline, padline in izip(textlines, padlines):
            yield padline
            yield textline
            yield ''

    def _cipher_padfill(self, ciphertext):
        self._ensure_pad_is_lines()
        padlines = [line.strip() for line in self.pad if line.strip()]
        cipherlines = _chunk_chars_into_lines(ciphertext, self._nullchar,
                                              _get_textwidth(padlines))
        for cipherline, padline in izip(cipherlines, padlines):
            yield padline
            yield ''
            yield cipherline

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

    def _cipher_textfill(self, padfilled_cipher_lines):
        for i, line in enumerate(padfilled_cipher_lines):
            lineno = i + 1
            if self._is_cipherline(lineno):
                padline = padfilled_cipher_lines[i - abs(self._padline_offset)]
                textline = \
                    padfilled_cipher_lines[i - abs(self._textline_offset)]
                yield padline
                yield self._textline_from_cipherline_and_padline(line, padline)
                yield line

    def _cipherline_from_padline_and_textline(self, padline, textline):
        ret = []
        for padchar, textchar in izip(padline, textline):
            if textchar == self._nullchar:
                ret.append(self._nullchar)
                continue
            charnum = self._as_nums[padchar] + self._as_nums[textchar]
            idx = charnum if charnum <= self._cipher_ceil else \
                    charnum % self._cipher_ceil
            ret.append(self._as_alpha[idx])
        return ''.join(ret)

    def _textline_from_cipherline_and_padline(self, cipherline, padline):
        ret = []
        for ciphercar, padchar in izip(cipherline, padline):
            if ciphercar == self._nullchar:
                ret.append(self._nullchar)
                continue
            charnum = self._as_nums[ciphercar] - self._as_nums[padchar]
            idx = charnum if charnum <= self._cipher_ceil else \
                    charnum % self._cipher_ceil
            if idx < 0:
                idx = self._cipher_ceil + idx
            ret.append(self._as_alpha[idx])
        return ''.join(ret)

    def _is_padline(self, lineno):
        return self._is_cipherline(lineno + abs(self._padline_offset))

    def _is_textline(self, lineno):
        return self._is_cipherline(lineno + abs(self._textline_offset))

    def _is_cipherline(self, lineno):
        return not lineno % self._pad_modulo


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


def _as_line_chunks(chars, nullchar, 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 _get_textwidth(textlines):
    return max([len(line) for line in textlines])


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


class _FakeDevUrandom(object):

    @staticmethod
    def read(nchars):
        import random
        ret = []
        for i in range(0, nchars):
            ret.append(chr(random.randint(0, 255)))
        return ''.join(ret)


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


_DEV_URANDOM = {}