Data whitening (scrambling) with 9-bit LFSR

whitening.py

import sys
import os
import numpy
import random
import string
import matplotlib.pyplot as plt

#data = numpy.random.bytes(10000000)
data = [random.choice(string.printable) for x in xrange(1000000)]
#data = numpy.random.bytes(100)

#data = [5, 10, 15, 20, 100, 150, 200, 255, 0, 128, 177, 212, 200, 78, 79, 44, 57, 92, 7,12,123,211,109,120]
#data = [1,2,3,4, 0x1F, 0xFF, 0xA4]
#data = [0x1F, 0xFF, 0xA4] * 1000
#data = [0x0f, 0x00, 0x01, 0x02]
#data = [chr(x) for x in data]

def reversebits(word, numbits=8):
    return sum(1<<(numbits-1-i) for i in range(numbits) if word>>i&1)

def lfsr_generator(count=8):
    word = 0x1FF
    while True:
        yield word
        for i in range(count):
            bit = (word >> 5) ^ (word >> 0) & 1
            outword = ((word >> 1) & 0xFF) | (bit << 8)
            word = outword & 0x1FF

def lfsr(word, count=8):
    word = word & 0x1FF
    for i in range(count):
        bit = (word >> 5) ^ (word >> 0) & 1
        outword = ((word >> 1) & 0xFF) | (bit << 8)
        word = outword & 0x1FF
    return word


def testlfsr():
    x = 0x1FF
    for i in range(10):
        print "%d: 0x%.4x\t%s" % (i, x, bin(x))
        x = lfsr(x)

def lfsr_list(list, reverse=False, display=False):
    xordata = []
    xorval = 0x1FF
    for x in list:
        xorbyte = chr(ord(x) ^ (xorval & 0xFF))
        if display:
            print "0x%.2X ^ 0x%.2X == 0x%.2X" % (ord(x), xorval&0xFF, ord(xorbyte))
        xordata.append(xorbyte)
        xorval = lfsr(xorval)
    return xordata
    
def check_restore(list):
    print "Restore data"
    restored = lfsr_list(list, True)
    for idx in xrange(len(data)):
        if ord(data[idx]) != ord(restored[idx]):
            print "Mismatch on byte %d!  0x%.2x != 0x%.2x" % (idx, ord(data[idx]), ord(restored[idx]))
            sys.exit(1)
    print "Double XORed list okay."

def drawWhiteningPlots(millions=1, randfunc=random.gauss, *randargs):
    if not randargs:
        randargs = (0.5,.1)
    data = [int(randfunc(*randargs)*255)%256 for i in xrange(int(1000000*millions))]
    chrdata = [chr(x) for x in data]
    xdata = lfsr_list(chrdata)
    xdata = [ord(x) for x in xdata]

    plt.subplot(2,1,1)
    plt.hist(data, bins=256)
    plt.xlabel('Byte')
    plt.ylabel('Count')
    plt.title("%.1f Million Random Bytes"%millions)

    plt.subplot(2,1,2)
    plt.hist(xdata, bins=256)
    plt.xlabel('Byte')
    plt.ylabel('Count')
    plt.title("%.1f Million Whitened Bytes"%millions)    
    
def maxBitRuns(list):
    maxone = 0
    sumone = 0
    countone = 0
    
    maxzero = 0
    sumzero = 0
    countzero = 0
    
    bitcount = 0
    lastbit = 0
    for idx,char in enumerate(list):
        byte = ord(char)
        for i in range(7,-1,-1):
            bit = byte & (1<<i) != 0
            if bit != lastbit:
                if (bitcount > 200):
                    print "WARNING: long bitcount at index %d (%d bits)" % (idx, bitcount)
                    print "%s" % [bin(ord(x)) for x in list[idx - bitcount/8 - 1:idx+1]]
                    print "%s" % [hex(ord(x)) for x in list[idx - bitcount/8 - 1:idx+1]]
                if lastbit and bitcount > maxone:
                    maxone = bitcount
                elif not lastbit and bitcount > maxzero:
                    maxzero = bitcount

                if lastbit:
                    sumone += bitcount
                    countone += 1
                else:
                    sumzero += bitcount
                    countzero += 1
                    
                bitcount = 1
            else:
                bitcount += 1
            lastbit = bit
    if lastbit and bitcount > maxone:
        maxone = bitcount
    elif not lastbit and bitcount > maxzero:
        maxzero = bitcount
    return (maxzero,maxone, float(sumzero)/countzero, float(sumone)/countone)



if __name__ == "__main__":
    print "XOR data"
    xordata = lfsr_list(data, display=False)
    
    #check_restore(xordata)
    
    print "Max bit runs in %d bytes" % len(data)
    maxbits = maxBitRuns(data)
    print "Max: %s" % str(maxbits)
    
    print "Max bit runs in whitened data"
    maxbits = maxBitRuns(xordata)
    print "Max: %s" % str(maxbits)

    plt.figure(1)
    drawWhiteningPlots(0.25, random.gauss, 0.1,0.2)
    plt.figure(2)
    drawWhiteningPlots(0.25, random.expovariate, 1.0/0.5)
    plt.figure(3)
    drawWhiteningPlots(0.25, random.triangular, 0,1.0)
    plt.figure(4)
    drawWhiteningPlots(0.25, random.gauss, 0.5, 0.1)
    plt.figure(5)
    drawWhiteningPlots(0.25, random.lognormvariate, 0.3, 0.2)

    plt.show()