# for python 3.0
# from 李勃
import struct,sys,os,binascii
"""
  RC4加密算法
  16*16 S盒
  加密單元:short
"""
def RC4(pkey,keylen,pin,dlen):
  N=65536
  S = list(range(N))
  j = 0
  for i in range(N):
    j = (j + S[i] + pkey[i%keylen])%N
    temp = S[i]
    S[i] = S[j]
    S[j] = temp
  i = j = 0
  pout= b''
  for x in range(dlen):
    i = i+1
    j = (j + S[i])%N
    temp = S[i]
    S[i] = S[j]
    S[j] = temp
    pout += struct.pack('H',pin[x]^S[(S[i]+S[j])%N])
  return(pout)
# bytes->short
def Coding(data):
  if(len(data)%2):
    data+=b'\0'
  dlen = len(data)//2
  return(struct.unpack(str(dlen)+'H',data))
# short->bytes
def unCoding(data):
  d=b''
  for i in range(len(data)):
    d += struct.pack('H',data[i])
  return(d)
#產生32字節(jié)密鑰
def CreatKey(Keyt):
  pl = len(Keyt)
  Key=b''
  r=0
  for i in range(32):
    k=(Keyt[r%pl]+i)%256
    Key+= struct.pack('B',k)
    r+=1
  return Key
#更新密鑰
def UpdataKey(Keyt):
  Key = unCoding(Keyt)
  #循環(huán)左移
  Key = Key[1:] + struct.pack('B',Key[0])
  tem=0
  #求和
  for i in range(len(Key)):
    tem += Key[i];
  Keyo=b''
  #Xor
  for i in range(len(Key)):
    Keyo += struct.pack('B',(Key[i]^tem)%256)
    tem += Keyo[i]>>3
    tem = tem % 256
  return(Coding(Keyo))
if __name__ == '__main__':
  #獲得輸入文件
  if len(sys.argv)==1:
    filename = input('源文件: ')
  else:
    filename = sys.argv[1]
   
  try:
    fin = open(filename,'rb')
  except:
    print('打開文件失敗！')
    input()
    sys.exit()
  print(filename)
  #打開輸出文件
  if filename[-4:]=='.RC4':
    eID = 1
    key=input('輸入解密密鑰: ').encode()
    ofilename = filename[:-4]
  else:
    eID = 2
    key=input('輸入加密密鑰: ').encode()
    ofilename = filename+'.RC4'
  key = Coding(CreatKey(key))
  key = UpdataKey(key)
 
  #處理重名
  while os.path.exists(ofilename):
    ofilename = os.path.dirname(ofilename)+ '\\副本 '+ os.path.basename(ofilename)
  fout = open(ofilename,'wb')
  print(ofilename)
  #解密
  if eID==1:
    #讀文件長度
    filelen = struct.unpack('I',fin.read(4))[0]
    print('FlieLen =',filelen,'\n......')
    while 1:
      #讀塊大小
      ps= fin.read(2)
      if not ps:
        #文件結束
        break
      packsize = struct.unpack('H',ps)[0]
      #讀數(shù)據(jù)
      dd=fin.read(packsize)
      #解密
      dd=Coding(dd)
      x = RC4(key,len(key),dd,len(dd))
      key = UpdataKey(key)
      #crc
      crc = struct.unpack('I',fin.read(4))[0]
      if binascii.crc32(x)!=crc:
        print('CRC32校驗錯誤！',crc,binascii.crc32(x))
        input()
        sys.exit()
      fout.write(x)
    #裁剪末尾填充位
    fout.truncate(filelen)
  #加密
  elif eID==2:
    #獲得文件長度
    fin.seek(0,2)
    filelen = fin.tell()
    print('FlieLen =',filelen,'\n......')
    fin.seek(0,0)
    fout.write(struct.pack('I',filelen))
    while 1:
      #讀數(shù)據(jù)
      dd=fin.read(65534)
      if not dd:
        #文件結束
        break
      #末尾填充
      srl = len(dd)
      if srl%2:
        srl+=1;
        dd+=b'\0'
      #crc
      crc = struct.pack('I',binascii.crc32(dd))
      #加密數(shù)據(jù)
      dd=Coding(dd)
      x = RC4(key,len(key),dd,len(dd))
      key = UpdataKey(key)
      #寫入文件
      fout.write(struct.pack('H',srl))
      fout.write(x)
      fout.write(crc)
  fin.close()
  fout.close()
  print('OK!')
  input()