BaldDecrypt.py

import itertools


# Re-implementation of important methods from the VB code
init_step_count = 1024

def step_waves(waves):
  for i in xrange(len(waves)):
    if i + 1 < len(waves):
      waves[i] += waves[i+1] - 8
    
    # This is almost modulo 15, but not quite
    if waves[i] > 15:
      waves[i] -= 15
    elif waves[i] < 0:
      waves[i] += 15
  
  return waves[0]

def setup_waves(key):
  waves = [ int(key[i], 16) for i in xrange(len(key)) ]
  for _ in xrange(init_step_count):
    step_waves(waves)
  return waves

def crypt(key, message):
  output = []
  waves = setup_waves(key)
  for m in message:
    hi = (ord(m) / 16) ^ step_waves(waves)
    lo = (ord(m) % 16) ^ step_waves(waves)
    output.append(16 * hi + lo)
  return ''.join(chr(o) for o in output)


# Tests
# Actually, bad tests, since these didn't catch my modulo operator
# mistake in step_waves(). Oops.
key = '010203040506'
waves = [ 12, 1, 13, 10, 0, 13, 13, 0, 2, 6, 7, 6 ]
assert setup_waves(key) == waves


# Codebreaking routines

def extract_stream(ct, pt, keylen):
  # Extract the stream at the end of the ciphertext, from which we
  # should be able to derive the waves
  stream = []
  for c, p in zip(ct, pt):
    x = ord(c) ^ ord(p)
    hi, lo = x / 16, x % 16
    stream.append(hi)
    stream.append(lo)
  return stream

def display_matrix(m):
  dm = []
  for r in xrange(len(m)):
    dm.append([])
    for c in xrange(len(m[r])):
      if m[r][c] is None:
        dm[r].append('-')
      else:
        dm[r].append(','.join(str(x) for x in m[r][c]))
  
  lens = [ max(len(dm[r][c]) for r in xrange(len(m))) \
    for c in xrange(len(m[0])) ]
  
  for row in dm:
    for c, val in enumerate(row):
      print val.rjust(lens[c]),
    print
  

def uncrypt(ct, pt, keylen):
  # Compute the stream at the end of the ciphertext
  stream = extract_stream(ct, pt, keylen)
  
  # Consider a matrix where each row represents 'waves' at one step during 
  # encipherment. For a small key like '1234', it might look like this:
  #
  #     1  2  3  4
  #    10 12 14  4
  #    14  3 10  4
  #     9  5  6  4
  #     6  3  2  4
  #     1 12 13  4
  #     5  2  9  4
  #    14  3  5  4
  #     9  0  1  4
  #
  # Note the first column represents the key stream because waves[0] is what is
  # emitted each time.
  #
  # Suppose now we wanted to reconstruct this matrix, and all we had was the
  # first column:
  #
  #     1  -  -  -
  #    10  -  -  -
  #    14  -  -  -
  #     9  -  -  -
  #     6  -  -  -
  #     1  -  -  -
  #     5  -  -  -
  #    14  a  -  -
  #     9  -  -  -
  #
  # Look at the bottom-left corner. Since we know that step_waves computes a
  # function f(14, a) = 9, we can solve this function for `a`. We would find
  # a = 3. We can solve the second column, except for the last row, in this way.
  #
  #     1  2  -  -
  #    10 12  -  -
  #    14  3  -  -
  #     9  5  -  -
  #     6  3  -  -
  #     1 12  b  -
  #     5  2  -  -
  #    14  3  -  -
  #     9  -  -  -
  #
  # And now we can repeat with the third column to solve f(12, b) = 2, etc. At
  # the end, we have:
  #
  #    10 12 14  4
  #    14  3 10  4
  #     9  5  6  4
  #     6  3  2  4
  #     1 12 13  4
  #     5  2  9  -
  #    14  3  -  -
  #     9  -  -  -
  #
  # Here, the first row is the waves state that produced the first byte of the
  # key stream.
  
  # Compute a map that lets us quickly solve f(a, x) = b for x. Note that there
  # is not always a unique solution. For example f(9, 0) = 1 and f(9, 15) = 1.
  # So we represent them as a tuple.
  _rev = {}
  for a in xrange(16):
    for b in xrange(16):
      sol = []
      for x in xrange(16):
        waves = [a, x]
        if step_waves(waves) == b:
          sol.append(x)
      _rev[a,b] = tuple(sol)
  
  def lookup(a, b, where):
    sol = set()
    for aa in a:
      for bb in b:
        sol.update(where[aa, bb])
    return tuple(sol)
  
  def lookup_rev(a, b):
    return lookup(a, b, _rev)
  
  # And a map that lets us solve f(a, b) = x. This is always unique.
  _fwd = {}
  for a in xrange(16):
    for b in xrange(16):
      waves = [a, b]
      _fwd[a, b] = (step_waves(waves),)
      
  def lookup_fwd(a, b):
    return lookup(a, b, _fwd)
  
  # Construct our matrix
  ncols = keylen
  nrows = len(stream)
  m = [ [None] * ncols for _ in xrange(nrows) ]
  
  # Fill in the stream in the first column
  for r in xrange(nrows):
    m[r][0] = (stream[r],)
  
  # Fill in the rest of the columns
  for c in xrange(1, ncols):
    for r in xrange(nrows - c):
      for v in xrange(15):
        m[r][c] = lookup_rev(m[r][c-1], m[r+1][c-1])
  
  # Now for any cell that contains multiple values, we see if we can eliminate
  # all but one of them
  for r in xrange(1, nrows):
    for c in xrange(ncols - 1):
      if m[r][c] is None:
        continue
      m[r][c] = lookup_fwd(m[r-1][c], m[r-1][c+1])
  
  #display_matrix(m)
  
  
  # Next we need to undo the initial permutations to find a previous row,
  # (a b c d) such that f(a, b) = 10, f(b, c) = 12, etc.
  #
  #     a  b  c  d
  #    10 12 14  4
  #
  # Here `d` is unbounded, so we have to consider all possible options. But,
  # given a choice for `d`, this forces `c`, and so on.
  #
  # But, note that the last digit of the wave is "stuck"! It is never modified
  # by step_waves(). So, we know that d = 4.
  #
  # We can repeat this as many times as we want to compute earlier rows.
  
  waves = m[0]
  del m
  
  for _ in xrange(init_step_count + 1):
    row = [None] * (ncols-1) + [waves[-1]]
    for c in xrange(ncols-2, -1, -1):
      row[c] = lookup_rev(row[c+1], waves[c])
    waves = row
  
  # Now we just enumerate all possible keys (not sure there are multiples)
  keys = set()
  for key in itertools.product(*waves):
    keys.add(''.join('0123456789ABCDEF'[k] for k in key))
  return list(keys)


# Test that we find our key
pt = open('img.jpg').read()
ct = open('imgenc.jpg').read()
for key in uncrypt(ct, pt, 512/8):
   assert crypt(key, pt) == ct
   print key, 'works!'