Format preserving encryption using a Feistel network

fpe_feistel.py

""" Format preserving encryption using a Feistel network
    See https://stackoverflow.com/a/51429458/4014959
    Written by PM 2Ring 2016.08.22

    Version using Blake2 hashing 2021.06.14
    Added pair permutation 2024.01.17
"""

from hashlib import blake2s as blake
from random import Random

class FormatPreserving:
    """ Invertible permutation of integers in range(stop), 0 < stop <= 2**64
        using a simple Feistel network.
    """
    def __init__(self, stop, keystring):
        if not 0 < stop <= 1 << 64:
            raise ValueError('stop must be <=', 1 << 64)

        # The highest number in the range
        self.maxn = stop - 1

        # Get the number of bits in each part by rounding
        # the bit length up to the nearest even number
        self.shiftbits = -(-self.maxn.bit_length() // 2)
        self.lowmask = (1 << self.shiftbits) - 1
        self.lowshift = 32 - self.shiftbits

        # Make 4 32 bit round keys from the keystring.
        # Create an independent random stream so we
        # don't intefere with the default stream.
        stream = Random()
        stream.seed(keystring)
        self.keys = [stream.getrandbits(32).to_bytes(4, "little") for _ in range(4)]
        self.ikeys = self.keys[::-1]

    def feistel(self, n, keys):
        # Split the bits of n into 2 parts & perform the Feistel
        # transformation on them.
        left, right = n >> self.shiftbits, n & self.lowmask
        for key in keys:
            h = blake(right.to_bytes(4, "little"), key=key, digest_size=4).digest()
            left, right = right, left ^ (int.from_bytes(h, "little") >> self.lowshift)
        return (right << self.shiftbits) | left

    def fpe(self, n, reverse=False):
        keys = self.ikeys if reverse else self.keys
        while True:
            # Cycle walk, if necessary, to ensure n is in range.
            n = self.feistel(n, keys)
            if n <= self.maxn:
                return n

def test():
    print('Shuffling a small number')
    maxn = 10
    fpe = FormatPreserving(maxn, 'secret key string')
    for i in range(maxn):
        a = fpe.fpe(i)
        b = fpe.fpe(a, reverse=True)
        print(i, a, b)

    print('Shuffling a small number, with a slightly different keystring')
    fpe = FormatPreserving(maxn, 'secret key string.')
    for i in range(maxn):
        a = fpe.fpe(i)
        b = fpe.fpe(a, reverse=True)
        print(i, a, b)

    print('Here are a few values for a large maxn')
    maxn = 10000000000000000000
    print('maxn =', maxn)
    fpe = FormatPreserving(maxn, 'secret key string')
    for i in range(10):
        a = fpe.fpe(i)
        b = fpe.fpe(a, reverse=True)
        print('{}: {:19} {}'.format(i, a, b))

    print('Using a set to test that there are no collisions...')
    maxn = 1 << 16 #100000
    print('maxn', maxn)
    fpe = FormatPreserving(maxn, 'secret key string')
    a = {fpe.fpe(i) for i in range(maxn)}
    print(len(a) == maxn)

    print('Testing that the operation is bijective...')
    for i in range(maxn):
        a = fpe.fpe(i)
        b = fpe.fpe(a, reverse=True)
        assert b == i, (i, a, b)
    print('yes')

#test()

# Permute (n, m) pairs
@interact
def test(n=3, m=5, seed="randseed", auto_update=False):
    nm = n * m
    fpe = FormatPreserving(nm, seed)
    for i in range(nm):
        j = fpe.fpe(i)
        print(f"{i:3}: {divmod(j, m)}")

Live version on SageMathCell

See https://en.wikipedia.org/wiki/Format-preserving_encryption