tools.py

from functools import partial
from StringIO import StringIO

import random

try:
    from IPython import embed as debug
except ImportError:
    pass

def egcd(a, b):
    """
    Computes the extended GCD for (a,b). That is, it computes integers x and y
    such that ax + by = gcd(a, b) as well as gcd(a, b).

    :param a: First parameter for GCD.
    :param b: Second parameter for GCD.
    :return:
    """
    if a == 0:
        return (b, 0, 1)
    else:
        g, y, x = egcd(b % a, a)
        return g, x - (b // a) * y, y

def modinv(a, m):
    """
    Computes the modular inverse of a mod m.
    http://stackoverflow.com/questions/4798654/modular-multiplicative-inverse-function-in-python

    :param a: a in a^-1 mod m
    :param m: m in a^-1 mod m
    :return: a^-1 mod m or None if no inverse exists
    """
    g, x, y = egcd(a, m)
    if g != 1:
        return None
    else:
        return x % m


def join(s):
    """
    Turns an array of strings into one string.

    :param s: Array of string to join.
    :return: One string made by connecting all strings in the array.
    """
    return ''.join(s)

def split(s, n=None):
    """
    Split a string into portions of length n. If n is not supplied the
    string is split in half. Some examples:

    .. testcode::

        from crypto.tools import split

        print split("ABCDEF", 1)
        print split("ABCDEF", 2)
        print split("ABCDEF")

    .. testoutput::

        ['A', 'B', 'C', 'D', 'E', 'F']
        ['AB', 'CD', 'EF']
        ['ABC', 'DEF']


    :param s: A string
    :param n: the length of string portions
    :return s[]: an array of the portions of s
    """
    if n is None:
        return [s[:len(s) / 2], s[len(s) / 2:]]
    else:
        return [l for l in iter(partial(StringIO(s).read, n), '')]


def string_to_int(s):
    """
    Converts s to an int

    :param s: A string
    :return: The integer representation of the string s.
    """
    x = 0
    for i in range(len(s)):
        x += ord(s[i]) << (len(s)-1 - i) * 8
    return x


def int_to_string(x, l=None):
    """
    Converts a number to a string with length l

    :param x: A number between 0 and 2 ** l
    :param l: The length of the string to be returned
    :return: string
    """
    s = ''
    while x != 0:
        char = chr(x & 0xFF)
        x >>= 8
        s = char + s

    if l is None:
        return s
    else:
        return "\x00" * (l - len(s)) + s


def add_int_to_string(s, num, block_size):
    """
    Adds a number to a string

    :param s: String of length block_size
    :param num: A number s.t. 0 <= num < block_size
    :return: A string
    """
    x = (string_to_int(s) + num) % (2 ** block_size)
    s = int_to_string(x, block_size/8)
    return s


def xor_strings(s1, s2):
    """
    Returns the bitwise XOR of s1 and s2. If len(s1) != len(s2) the resulting
    XOR operation will be the size of the bigger string.

    :param s1: first string in XOR
    :param s2: second string in XOR
    :return: result of bitwise XOR of s1 and s2.
    """

    if len(s1) < len(s2):
        s1 = s1 + ("\x00" * (len(s2) - len(s1)))
    elif len(s2) < len(s1):
        s2 = s2 + ("\x00" * (len(s1) - len(s2)))


    return "".join(chr(ord(x) ^ ord(y)) for x, y in zip(s1, s2))


def bitwise_complement_string(s):
    """
    Returns the bitwise complement of s.

    :param s: string obejct to complement
    :return: result of bitwuise complement
    """
    return xor_strings(s, len(s) * "\xFF")

# https://rosettacode.org/wiki/Miller%E2%80%93Rabin_primality_test#Python:_Proved_correct_up_to_large_N
def is_prime(n, _precision_for_huge_n=16):
    def _try_composite(a, d, n, s):
        if pow(a, d, n) == 1:
            return False
        for i in range(s):
            if pow(a, 2**i * d, n) == n-1:
                return False
        return True # n  is definitely composite

    if not hasattr(is_prime, '_known_primes'):
        is_prime._known_primes = [2, 3]
        is_prime._known_primes += [x for x in range(5, 1000, 2) if is_prime(x)]

    if n in is_prime._known_primes or n in (0, 1):
        return True
    if any((n % p) == 0 for p in is_prime._known_primes):
        return False
    d, s = n - 1, 0
    while not d % 2:
        d, s = d >> 1, s + 1
    # Returns exact according to http://primes.utm.edu/prove/prove2_3.html
    if n < 1373653:
        return not any(_try_composite(a, d, n, s) for a in (2, 3))
    if n < 25326001:
        return not any(_try_composite(a, d, n, s) for a in (2, 3, 5))
    if n < 118670087467:
        if n == 3215031751:
            return False
        return not any(_try_composite(a, d, n, s) for a in (2, 3, 5, 7))
    if n < 2152302898747:
        return not any(_try_composite(a, d, n, s) for a in (2, 3, 5, 7, 11))
    if n < 3474749660383:
        return not any(_try_composite(a, d, n, s) for a in (2, 3, 5, 7, 11, 13))
    if n < 341550071728321:
        return not any(_try_composite(a, d, n, s) for a in (2, 3, 5, 7, 11, 13, 17))
    # otherwise
    return not any(_try_composite(a, d, n, s)
                   for a in is_prime._known_primes[:_precision_for_huge_n])


def prime_between(s, e):
    candidate = random.randint(s, e)

    while not (is_prime(candidate)):
        candidate = random.randint(s, e)

    return candidate

def exp(a, n, N):
	r = 1
	for i in range(n.bit_length() - 1, -1, -1):
		r = ((r*r) * (a if (n >> i) & 0x1 else 1)) % N

	return r