# analysis of SHA-256 data and key bit propagation
#
# list and set elements here represent data or key bit indices
#
# each element, e.g. W[i][j] and the function argument x,
# is a list or set of bit indices
#
# Note: does not account for carry in addition formulas (W,T0,T1,T2,a,e)
#
# little-endian bit order is used, i.e. x[0] is the lsb bit 0,
# so big-endian ROTR and SHR are implemented by shifting left
#
# The message schedule W[t] is kept as a list of bits only up to t = 31
# due to excessive memory requirements beyond that.

# set KEY = True to analyze key bits instead of data bits
#
KEY = True
#-- KEY = False

# max number of rounds to consider: 16 <= N <= 64
#
N = 17

# circular rotate right by n bits
#
def ROTR(x,n):
  if type(x[0]) == set:
    R = [ set() for i in range(32)]
    for i in range(32):
      k = (i+n)%32; R[i] |= x[k]
  elif type(x[0]) == int:
    R = [ None for i in range(32)]
    for i in range(32):
      k = (i+n)%32; R[i] = x[k]
  else:
    R = [ [] for i in range(32)]
    for i in range(32):
      k = (i+n)%32; R[i].extend(x[k])
  return R

# shift right by n bits
#
def SHR(x,n):
  if type(x[0]) == set:
    R = [ set() for i in range(32)]
    for i in range(32-n):
      k = i+n; R[i] |= x[k]
  else:
    R = [ [] for i in range(32)]
    for i in range(32-n):
      k = i+n; R[i].extend(x[k])
  return R

# SHA-256 Ch function
#
def Ch(x,y,z):
  R = [ set() for i in range(32)]
  for i in range(32):
    R[i] = x[i] | y[i] | z[i]
  return R

# SHA-256 Maj function
#
def Maj(x,y,z):
  R = [ set() for i in range(32)]
  for i in range(32):
    R[i] = x[i] | y[i] | z[i]
  return R

# SHA-256 Sigma0 function
#
def Sigma0(x):
  w = ROTR(x,2); y = ROTR(x,13); z = ROTR(x,22)
  R = [ set() for i in range(32)]
  for i in range(32):
    R[i] = w[i] | y[i] | z[i]
  return R

# SHA-256 Sigma1 function
#
def Sigma1(x):
  w = ROTR(x,6); y = ROTR(x,11); z = ROTR(x,25)
  R = [ set() for i in range(32)]
  for i in range(32):
    R[i] = w[i] | y[i] | z[i]
  return R

# SHA-256 sigma0 function
#
def sigma0(x):
  w = ROTR(x,7); y = ROTR(x,18); z = SHR(x,3)
  if type(x[0]) == set:
    R = [ set() for i in range(32)]
    for i in range(32):
      R[i] = w[i] | y[i] | z[i]
  else:
    R = [ [] for i in range(32)]
    for i in range(32):
      R[i].extend(list(w[i]))
      R[i].extend(list(y[i]))
      R[i].extend(list(z[i]))
  return R

# SHA-256 sigma1 function
#
def sigma1(x):
  w = ROTR(x,17); y = ROTR(x,19); z = SHR(x,10)
  if type(x[0]) == set:
    R = [ set() for i in range(32)]
    for i in range(32):
      R[i] = w[i] | y[i] | z[i]
  else:
    R = [ [] for i in range(32)]
    for i in range(32):
      R[i].extend(list(w[i]))
      R[i].extend(list(y[i]))
      R[i].extend(list(z[i]))
  return R

# set Key and K, or Data and D
#
if KEY:

  # 256-bit key
  #
  Key = [i for i in range(256)]

  # 8 32-bit pieces of the Key
  #
  K = [ [] for i in range(8)]

  for i in range(8):
    K[7-i] = Key[i*32:(i+1)*32]

else:

  # 512-bit input data block
  #
  Data = [i for i in range(512)]

  # 16 32-bit pieces of the data
  #
  D = [ [] for i in range(16)]

  for i in range(16):
    D[15-i] = Data[i*32:(i+1)*32]

# message schedule
#
# W[i][j] will be a list of data or key bits affecting bit j of W[i]
# 
# WS is W using sets instead of lists
#
W  = [[  []   for j in range(32)] for i in range(N)]
WS = [[ set() for j in range(32)] for i in range(N)]

# W[0]...W[15]
#
if KEY:

  KeyIndex = [ 0, 1, 2, 3, 4, 5, 6, 7, 3, 4, 5, 6, 7, 1, 2, 0]

  for i in range(16):
    ki = KeyIndex[i]
    x = ROTR( K[ki], 3*ki+(i>>3))
    for j in range(32):
      W[i][j].append(x[j])
      WS[i][j].add(x[j])

else:

  for i in range(16):
    for j in range(32):
      W[i][j].append(D[i][j])
      WS[i][j].add(D[i][j])

# W[16]...
#
for i in range(16,min(N,32)):
  y = sigma1(W[i-2])
  z = sigma0(W[i-15])
  if KEY: ki = (63-i)%8; x = ROTR( K[ki], 3*ki+(i>>3))
  for j in range(32):
    W[i][j].extend(y[j])
    W[i][j].extend(W[i-7][j])
    W[i][j].extend(z[j])
    W[i][j].extend(W[i-16][j])
    if KEY: W[i][j].append(x[j])
    WS[i][j] = set(W[i][j])

#
# omit W[t] as lists for t >= 32 due to excessive memory requirements
#
for i in range(32,N):
  if KEY: ki = (63-i)%8; x = ROTR( K[ki], 3*ki+(i>>3))
  for j in range(32):
    y = sigma1(WS[i-2]); z = sigma0(WS[i-15])
    WS[i][j] = WS[i-7][j] | WS[i-16][j] | y[j] | z[j]
    if KEY: WS[i][j].add(x[j])

# SHA-256 on one block
#
# N < 64 represents a reduced-round SHA-256.
#
def SHA256(N=20):
  global a, b, c, d, e, f, g, h, hprev, T0, T1, T2
  a = [ set() for i in range(32)]
  b = [ set() for i in range(32)]
  c = [ set() for i in range(32)]
  d = [ set() for i in range(32)]
  e = [ set() for i in range(32)]
  f = [ set() for i in range(32)]
  g = [ set() for i in range(32)]
  h = [ set() for i in range(32)]
  for t in range(N):
    T0 = [ set() for i in range(32)]
    y = Sigma1(e); z = Ch(e,f,g)
    for i in range(32):
      T0[i] = h[i] | y[i] | z[i]
    T1 = [ set(WS[t][i]) for i in range(32)]
    for i in range(32): # T1 = T0 + W[t]
      T1[i] |= T0[i]
    T2 = Sigma0(a); y = Maj(a,b,c)
    for i in range(32):
      T2[i] |= y[i]
    hprev = h; h = g; g = f; f = e; e = d
    d = c; c = b; b = a; a = [ set() for i in range(32)]
    for i in range(32):
      e[i] |= T1[i]; a[i] = T1[i] | T2[i]
  R = []
  R.extend(list(h)); R.extend(list(g)); R.extend(list(f)); R.extend(list(e))
  R.extend(list(d)); R.extend(list(c)); R.extend(list(b)); R.extend(list(a))
  return R

# count number of data or key bits used in each bit of the hash
#
def Hcount(N):
  H = SHA256(N)
  L = [ len(H[i]) for i in range(256)]
  return L
#
# results, data bits:
#  N min max
# 16 314 445
# 17 346 477
# 18 378 502
# 19 410 512
# 20 442 512
# 21 474 512
# 22 497 512
# 23 509 512
# 24 512 512
# ...
#
# results, key bits:
#  N min max
# 16 255 256
# 17 256 256
# 18 256 256
# 19 256 256
# 20 256 256
# ...

# find elements in W[P][n] which only occur once
# (or more generally, occur an odd number of times)
# and which do not occur directly in hprev[n] = bit n of h for t=P-1,
# and do not occur in W[P][n+1] or hprev[n+1]
#
# so flipping that bit would flip T1[n],
# and if the carry differs, and that bit doesn't directly affect hprev[n+1],
# then T1[n+1] would also flip
#
n = 0

if not KEY and N <= 32:
  
  H = SHA256(N); P = N-1
  
  Y = []; X = [-1]; X.extend(sorted(W[P][n])); X.append(-1) # sentinals
  
  for i in range(1,len(X)-1):
    if X[i] != X[i-1] and X[i] != X[i+1]: Y.append(X[i])
  
  Z = hprev[n] | WS[P][n+1] | hprev[n+1]
  
  A = [i for i in Y if i not in Z]
  
  #         n = 0             n = 1
  # N = 16, A = [0]           [1]
  # N = 17, A = [42, 49, 51]  [43, 50, 52]
  # N = 18, A = [10, 17, 19]
  # N = 19, A = [34, 38, 52]
  # N = 20, A = [2, 6, 20]

# note whether W contains data bits or key bits
#
print( "KEY =", KEY)