class Bits(object):
  __slots__ = ['ival','__sz','mask']

  def __init__(self,v,size=None,bitorder=-1):
    self.ival = self.__sz = self.mask = 0
    if isinstance(v,Bits):
      self.ival = v.ival
      self.__sz = v.size
      self.mask = v.mask
    elif isinstance(v,int):
      self.ival = abs(v*int(1))
      if self.ival>0 and (size is None):
        self.size = self.ival.bit_length()
    elif isinstance(v,list):
      self.size = len(v)
      self.ival = 0
      for x in reversed(v):
        self.ival = (self.ival<<1)|(x&1)
    elif isinstance(v,bytes):
      self.load(v,bitorder)
    else:
        raise TypeError(v)
    if size!=None: self.size = size

  def load(self,v,bitorder=-1):
    bytestr = bytes(v)
    self.size = len(bytestr)*8
    if bitorder==-1:
      l = [reverse_byte(c) for c in bytestr]
    else:
      l = bytestr
    v = 0
    for o in reversed(l):
      v = (v<<8) | o
    self.ival = v

  def __len__(self):
    return self.size

  def bit(self,i):
    if 0 <= i < self.__sz:
      return (self.ival>>i)&0x1
    elif 0<= -i <= self.__sz:
      return (self.ival>>(self.__sz+i))&0x1
    else:
      raise IndexError

  def int(self,sign=1):
    if sign==-1 and self.bit(-1)==1:
      return -(self.ival^self.mask)-1
    return self.ival&self.mask

  def __int__(self):
    return self.int()

  def __index__(self):
    return self.int()

  @property
  def size(self):
    return self.__sz

  @size.setter
  def size(self,v):
    self.__sz = v
    self.mask = (1<<v)-1
    self.ival &= self.mask

  def __repr__(self):
    c = self.__class__
    l = self.__sz
    s = self.ival
    return '<%s instance with ival=%x (len=%d)>'%(c,s,l)

  # binary string representation, bit0 1st.
  def __str__(self):
    xval = ("%x"%(self.ival&self.mask)).zfill(self.__sz//4+1)
    s = [hextab_r[int(x,16)] for x in xval]
    s.reverse()
    return u''.join(s)[:self.__sz]

  # byte string representation, bit0 1st (crypto notation).
  def __bytes__(self):
    v = self.ival&self.mask
    i = 0
    s = []
    while i<self.__sz:
      s.append(reverse_byte(v&0xff))
      v = v>>8
      i += 8
    return bytes(s)

  def bytes(self):
    return self.__bytes__()

  def hex(self):
    return codecs.encode(self.__bytes__(),'hex')

  def split(self,subsize,bigend=False):
    l = []
    i = 0
    while i<self.__sz:
      l.append(self[i:i+subsize])
      i += subsize
    if bigend: l.reverse()
    return l

  def todots(self):
    return u'|%s|'%str(self).replace('0',' ').replace('1','.')

# Basic comparison method ('is' operator), falls back to integer comparison.
#------------------------------------------------------------------------------
  def __cmp__(self,a):
    if not isinstance(a,Bits): raise AttributeError
    if self.size != a.size: raise ValueError
    return cmp(self.ival,a.ival)

# Enhanced comparison methods ('==' and '!=' operators).
#------------------------------------------------------------------------------
  def __eq__(self,a):
    if isinstance(a,Bits): a=a.ival
    return (self.ival==a)
#------------------------------------------------------------------------------
  def __ne__(self,a):
    if isinstance(a,Bits): a=a.ival
    return (self.ival!=a)

#------------------------------------------------------------------------------
  def __neg__(self):
    return Bits((-self.ival) % self.mask, self.size)

# Iterator for the class. Enables 'for b in self' expressions.
#------------------------------------------------------------------------------
  def __iter__(self):
    for x in range(self.size):
      yield self.bit(x)

# getitem operator defines b[i], b[i:j] and b[list] which returns the requested
# bit values as a int (0,1) or a list of such ints.
#------------------------------------------------------------------------------
  def __getitem__(self,i):
    if isinstance(i,int):
      return Bits(self.bit(i),1)
    elif isinstance(i,slice):
      start,stop,step = i.indices(self.__sz)
      if step==1 and stop>=start:
        v = (self.ival&((1<<stop)-1))>>start
        return Bits(v,stop-start)
      else:
        return self[range(self.__sz)[i]]
    else:
      v = 0
      for x in reversed(i):
        v = (v<<1)|((self.ival>>x)&1)
      return Bits(v,len(i))

# setitem operator defines b[i], b[i:j] and b[list] which allow to affect new
# values to these bits, from another object, int value or a bit list.
#------------------------------------------------------------------------------
  def __setitem__(self,i,v):
    if isinstance(i,int):
      assert v in (0,1)
      if   0<= i< self.__sz   : p=i
      elif 0<=-i<(self.__sz+1): p=self.__sz+i
      else: raise IndexError
      if v==0: self.ival &= (self.mask^((0x1)<<p))
      if v==1: self.ival |= (0x1)<<p
    else:
      v = Bits(v)
      if isinstance(i,slice):
        start,stop,step = i.indices(self.__sz)
        if step==1 and stop>start:
            mask = self.mask^((1<<stop)-1)^((1<<start)-1)
            self.ival = (self.ival&mask)|(v.ival<<start)
            return
        r = range(start,stop,step)
      else:
        r = i
      assert len(r)==len(v)
      for j,b in zip(r,v):
        self[j] = b

# unary bitwise operators. The result is a new object which has same length.
#------------------------------------------------------------------------------
  def __lshift__(self,i):
    res = Bits(self)
    res.ival = (res.ival<<i)&res.mask
    return res
  def __rshift__(self,i):
    res = Bits(self)
    res.ival = (res.ival>>i)&res.mask
    return res
  def __invert__(self):
    res = Bits(self)
    res.ival = res.ival ^ res.mask
    return res

  def zeroextend(self,size):
    if size>self.size:
      self.size = size
    return self

  def signextend(self,size):
    if size>self.size:
      m = self.mask
      s = self[-1]
      self.size = size
      if s==1:
        m ^= self.mask
        self.ival |= m
    return self

  def extend(self,sign,size):
    return self.signextend(size) if sign is True else self.zeroextend(size)

# binary operators, rvalue and lvalue implementations.
# (Note that resulting object length is max length.
#------------------------------------------------------------------------------
  def __and__(self,rvalue):
    if not isinstance(rvalue,Bits):
      obj = Bits(rvalue)
    else:
      obj = rvalue
    if self.size > obj.size:
      res = Bits(self)
    else:
      res = Bits(obj)
    res.ival = ( self.ival & obj.ival )
    return res
  def __or__(self,rvalue):
    if not isinstance(rvalue,Bits):
      obj = Bits(rvalue)
    else:
      obj = rvalue
    if self.size > obj.size:
      res = Bits(self)
    else:
      res = Bits(obj)
    res.ival = ( self.ival | obj.ival )
    return res
  def __xor__(self,rvalue):
    if not isinstance(rvalue,Bits):
      obj = Bits(rvalue)
    else:
      obj = rvalue
    if self.size > obj.size:
      res = Bits(self)
    else:
      res = Bits(obj)
    res.ival = ( self.ival ^ obj.ival )
    return res
  def __add__(self,rvalue):
    if not isinstance(rvalue,Bits):
      obj = Bits(rvalue)
    else:
      obj = rvalue
    if self.size > obj.size:
      res = Bits(self)
    else:
      res = Bits(obj)
    res.ival = ( self.ival + obj.ival )&res.mask
    return res
  def __sub__(self,rvalue):
    if not isinstance(rvalue,Bits):
      obj = Bits(rvalue)
    else:
      obj = rvalue
    if self.size > obj.size:
      res = Bits(self)
    else:
      res = Bits(obj)
    res.ival = ( self.ival - obj.ival )&res.mask
    return res

  def __rand__(self,lvalue):
    return (self & lvalue)
  def __ror__(self,lvalue):
    return (self | lvalue)
  def __rxor__(self,lvalue):
    return (self ^ lvalue)
  def __radd__(self,lvalue):
    return (self + lvalue)
  def __rsub__(self,lvalue):
    return Bits(lvalue,self.size)-self

# operator // is used for concatenation:
  def __floordiv__(self,rvalue):
    if not isinstance(rvalue,Bits):
      obj = Bits(rvalue)
    else:
      obj = rvalue
    size = self.size+obj.size
    return Bits(self.ival | obj.ival<<self.size, size)

  def bitlist(self,dir=1):
    l = list(self)
    if dir==-1: l.reverse()
    return l

# hamming weight of the object (count of 1s).
#------------------------------------------------------------------------------
  def hw(self):
    return self.bitlist().count(1)

# hamming distance to another object of same length.
#------------------------------------------------------------------------------
  def hd(self,other):
    if not isinstance(other,Bits):
      obj = Bits(other)
    else:
      obj = other
    if self.size != obj.size: raise ValueError
    return (self^obj).hw() class TDEA(object):
    size = 64
    blocksize = 64

    def __init__(self,K1,K2=None,K3=None):
        if len(K1)>8:
            assert K2 is None
            assert K3 is None
            K1,K2,K3 = K1[:8],K1[8:],K1[16:]
            if K3==b'': K3=K1
        if K2 is None:
            assert K3 is None
            K2 = K1
        if K3 is None:
            K3 = K1
        self.E1 = DES(K1)
        self.E2 = DES(K2)
        self.E3 = DES(K3)

    def enc(self,M):
        return self.E3.enc(self.E2.dec(self.E1.enc(M)))

    def dec(self,C):
        return self.E1.dec(self.E2.enc(self.E3.dec(C))) class DES(object):
    size = 64
    blocksize = 64

    def __init__(self,K):
        assert len(K)==self.size//8
        self.K = Bits(K,self.size)

    def enc(self,M):
        M = Bits(M)
        assert M.size==self.blocksize
        K = self.K
        k = PC1(K)
        blk = IP(M)
        L = blk[0:32]
        R = blk[32:64]
        for r in range(16):
            fout = F(R,k,r)
            L = L^fout
            L,R = R,L
        L,R = R,L
        C = Bits(0,64)
        C[0:32] = L
        C[32:64] = R
        return (IPinv(C)).bytes()

    def dec(self,C):
        C = Bits(C)
        assert C.size==self.blocksize
        K = self.K
        k = PC1(K)
        blk = IP(C)
        L = blk[0:32]
        R = blk[32:64]
        for r in reversed(range(16)):
            fout = F(R,k,r)
            L = L^fout
            L,R = R,L
        L,R = R,L
        M = Bits(0,64)
        M[0:32] = L
        M[32:64] = R
        return (IPinv(M)).bytes() def table_rKT(r,K):
    rks = table_rKS(r,K)
    rkt = []
    for n in range(12):
        rkt.append(list(range(256)))
    for v in range(256):
        re = Bits(v,8)
        for n in range(8):
            x = Bits(rks[n][re[0:6].ival],4)//re[(0,5,6,7)]
            rkt[n][re.ival] = x.ival
    for n in range(12):
        rkt[n] = tuple(rkt[n])
    return rks,tuple(rkt) def table_M1():
    l,r = range(32),Poly(list(range(32,64)))
    re = E(r).ival
    rbits = r.ival
    blk = []
    for b in range(12):
        blk.append([0]*8)
        if b<8:
            blk[b][0:6]=re[0:6]
            rbits.remove(re[0])
            rbits.remove(re[5])
            re = re[6:]
            blk[b][6:8]=l[0:2]
            l = l[2:]
        else:
            blk[b][0:4]=l[0:4]
            blk[b][4:8]=rbits[0:4]
            l = l[4:]
            rbits = rbits[4:]
    assert len(rbits)==0
    assert len(l)==0
    table = []
    for b in range(12):
        table.extend(blk[b])
    M = Poly(list(range(64)))
    return IP(M)[table].ival def table_M2():
    Mat = [Bits(0,96) for i in range(96)]
    SR,L,R = SRLRformat()
    newL = R.ival
    newR = list(zip(P(SR),L))
    RE = [newR[i] for i in E(Poly(list(range(len(L)))))]
    m=[]
    for r in range(8):
        m += RE[0:6]+newL[0:2]
        RE = RE[6:]
        newL  = newL[2:]
    rbits = getrbits_T_in()[16:]
    for r in range(4):
        m += newL[0:4]
        for b in rbits[:4]:
            m += [newR[b]]
        newL = newL[4:]
        rbits = rbits[4:]
    assert len(m)==96
    for v in range(96):
        l = m[v]
        try:
            for x in l: Mat[v][x]=1
        except TypeError:
            Mat[v][l]=1
        Mat[v] = Mat[v].ival
    return Mat,m def table_M3():
    ER,L,R = ERLRformat()
    C = Poly(R.ival+L.ival)
    return IPinv(C).ival class WhiteDES(object):

    def __init__(self,KT,tM1,tM2,tM3):
        self.KT = KT
        self.tM1 = tM1
        self.tM2 = tM2
        self.tM3 = tM3
        self.size = 64
        self.blocksize = 64

    def __FX(self,v):
        res = Bits(0,96)
        for b in range(96):
            res[b] = (v&self.tM2[b]).hw()%2
        return res

    def enc(self,M):
        assert len(M)==8
        M = Bits(M)
        blk = M[self.tM1]
        for r in range(16):
            t = 0
            for n in range(12):
                nt = t+8
                blk[t:nt] = self.KT[r][n][blk[t:nt]]
                t = nt
            blk = self.__FX(blk)
        return (blk[self.tM3]).bytes()

    def dec(self,C):
        assert len(C)==8
        #C = Bits(C)
        raise NotImplementedError class ECB(Mode):
    def __init__(self,cipher,pad=pkcs7):
        super().__init__(cipher,pad)
    # encryption mode
    def enc(self,M):
        C = []
        for b in self.iterblocks(M):
            C.append(self._cipher.enc(b))
        return b''.join(C)
    # decryption mode
    def dec(self,C):
        n,p = divmod(len(C),self.len)
        assert p==0
        P = BytesIO(C)
        M = []
        for b in range(n):
            M.append(self._cipher.dec(P.read(self.len)))
        return self.pad.remove(b''.join(M)) class CBC(Mode):
    def __init__(self,cipher,IV,pad=pkcs7):
        super().__init__(cipher,pad)
        assert len(IV)==self.len
        self.IV = IV
    # encryption mode
    def enc(self,M):
        C = [self.IV]
        for b in self.iterblocks(M):
            x = self.xorstr(b,C[-1])
            C.append(self._cipher.enc(x))
        return b''.join(C)
    # decryption mode
    def dec(self,C):
        l = self.len
        n,p = divmod(len(C),l)
        assert p==0
        M = []
        while len(C)>l:
            c = C[-l:]
            C = C[:-l]
            M.insert(0,self.xorstr(C[-l:],self._cipher.dec(c)))
        return self.pad.remove(b''.join(M))
import pytest
import codecs
from crysp.bits import Bits
from crysp.des import TDEA,DES
from crysp.wb import table_rKT,table_M1,table_M2,table_M3,WhiteDES
from crysp.mode import ECB,CBC

def test_des_wb():
    K = codecs.decode(b"0123456789ABCDEF",'hex')
    bK = Bits(K,64)
    b1 = b"Now is t"
    b2 = b"he time "
    b3 = b"for"+b"\5"*5
    KT = []
    for r in range(16):
        s,t = table_rKT(r,bK)
        KT.append(t)
    E  = DES(K)

    WT = ECB(WhiteDES(KT,table_M1(),table_M2()[0],table_M3()))