#
# Visualisation routine made by Karl-Uwe Frank
# 
# This is free and unencumbered software released into the public domain.
# 
# Anyone is free to copy, modify, publish, use, compile, sell, or
# distribute this software, either in source code form or as a compiled
# binary, for any purpose, commercial or non-commercial, and by any
# means.
# 
# In jurisdictions that recognize copyright laws, the author or authors
# of this software dedicate any and all copyright interest in the
# software to the public domain. We make this dedication for the benefit
# of the public at large and to the detriment of our heirs and
# successors. We intend this dedication to be an overt act of
# relinquishment in perpetuity of all present and future rights to this
# software under copyright law.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
# 
# For more information, please refer to <http://unlicense.org/>
# 


# files=$(ls -1 examples/byte_in_homebox_rc4-128-256-134217728-{aaa,bbb,ccc,ddd}.csv); python ./Plot-Byte_at_Homebox.py $files
# files=$(ls -1 examples/byte_in_homebox_rc4-{64,128,256}-256-134217728-abc.csv); python ./Plot-Byte_at_Homebox.py $files

import sys, os, math
import matplotlib.pyplot as plt
import numpy as np


#-----------------------------------------
# Simulate missing ANSI-C "printf"
#
def printf(*args):
  sys.stdout.write(*args)
  sys.stdout.flush()

def print_stderr(*args):
  sys.stderr.write(*args)
  sys.stderr.flush()


#-----------------------------------------
#
def ShowUsage(ThisName):
  print_stderr ( "\nUsage   : python %s CSV-file (max. 4 CSV files as argument)\n" % ThisName);
  print_stderr ( "\nExample : python %s byte_in_homebox_rc4-128-256-127926272-aaa.csv\n\n" % ThisName);


# -----------------------------------------
# Main 
#
def main():
  b         = [0] * 256
  min       = int(0)
  max       = int(0)
  files     = []
  ciphers   = []
  key_len   = [0] * 4
  byte_len  = int(0)
  key_count = int(0)
  init_val  = ''
#  colours   = ['blue','red','green','#DF71E7']
  colours   = ['blue','red','green','yellow']
  figTitle  = 'Visualise the Keystream of RC4-like Stream Cipher\n'
  xLabel    = 'Probability of byte occurrence in keystream at home box'
  yLabel    = 'Amount of appearance\n'
  algotitle = '' 
	
  #
  # Check if a Prarmeter is passed through
  PrgName = sys.argv[0]

  if (len(sys.argv) < 2 or len(sys.argv) > 5):
    ShowUsage(PrgName)
    sys.exit()

  for i in range(1,len(sys.argv)):
    ciphers.append('')
    files.append('')

    files[i-1] = sys.argv[i]
  
    temp = files[i-1].split('_')[3]
    ciphers[i-1] = temp.split('-')[0]
    key_len[i-1] = int(temp.split('-')[1])
    algotitle = algotitle +"_"+ciphers[i-1]

  fileName = os.path.basename(temp).split('-')[0:4]
  fileName = "figure_("+ algotitle +")-"+fileName[1]+"-"+fileName[2]+"-"+fileName[3]

  print fileName
     
  byte_len  = int(temp.split('-')[2])
  key_count = int(temp.split('-')[3])
  init_val  = (temp.split('-')[4]).split('.')[0]

  figTitle  = 'Visualise the first 256 Byte of a RC4-like Stream Cipher Keystream\ngenerated from ' + '{:,d}'.format(key_count) + ' random Keys'

  if len(files) > 1:
    figTitle  += ' for each Line'

  figTitle  += '\n'


  fig = plt.figure(figsize=(14,7), facecolor='white', edgecolor='black')
#  sub = fig.add_subplot(1, 1, 1)
  sub = fig.add_subplot(1, 1, 1, axisbg='#C7C7C7')

  plt.subplots_adjust(top=0.85)
  plt.subplots_adjust(left=0.10)
  plt.subplots_adjust(right=0.95)
  plt.subplots_adjust(bottom=0.2)

  
  for j in range(len(files)):
    with open(files[j]) as f:
        data = f.read()
        f.close()
    
    data = data.split('\n')
    
    for i in range(0,256):
      b[i] = int(data[i].split(',')[1])
      if b[i] < min or min == 0: 
        min = b[i]
      if b[i] > max: 
        max = b[i]
    

    sub.plot(b, colours[j], label=ciphers[j] +' ('+ str(key_len[j]) +' bit keys)')
 	
 
  sub.legend(loc='upper left', borderpad=0.5, labelspacing=0.5, borderaxespad=0.5, columnspacing=0.5, fontsize=10)

  xLabel += '\n{:,d}'.format(min) + ' min\n{:,d}'.format(max) +' max\n'
  
  if len(sys.argv) < 3:
    xLabel += '{:,d}'.format(max-min) +' spread\n'
  else:
    xLabel += '{:,d}'.format(max-min) +' over all spread\n'
  
  plt.title(figTitle, fontsize = 14)
  plt.ylabel(yLabel,  fontsize = 12)
  plt.xlabel(xLabel,  fontsize = 12)
  
  # major ticks every 8, minor ticks every 8                                      
  major_ticks = np.arange(0, 256, 8)                                              
  minor_ticks = np.arange(0, 256, 4)                                               
  
  sub.set_xticks(major_ticks)                                                       
  sub.set_xticks(minor_ticks, minor=True)                                           
  
  # set intensity of lines
  plt.grid(which='minor', alpha=0.2)                                                
  plt.grid(which='major', alpha=0.5)                                                

  plt.xticks([w*8 for w in range(0,256)])
  plt.xticks(rotation='vertical')
  
  plt.autoscale(tight=False)
  plt.grid(True)
  
  # fileName = "~/Desktop/myPlot.png"
  # plt.savefig(PNGfileName, format="png")
  plt.savefig(fileName +'.png', format="png")

  plt.show()


if __name__ == "__main__":
  main()
  sys.exit()