#!/usr/bin/env python

# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published by the
# Free Software Foundation, either version 3 of the License, or (at your
# option) any later version.
# 
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
# more details.
# 
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see .

# Inspired by Jeff Connelly
# Matthieu Weber, 2011

import time, sha, sys, socket
from optparse import OptionParser

parser = OptionParser()

parser.add_option("-l", 
    help="How much data should be output (use 'b' suffix for bits, 'K' for kByte, 'M' for MByte)", 
    type="string", 
    dest="data_length")

parser.add_option("-o",
    metavar="FILE", 
    help="Output file (default is the standard output)", 
    action="store", 
    type="string", 
    dest="output_file")

parser.add_option("-x", 
    help="Output hexadecimal numbers instead of raw data", 
    action="store_true", 
    dest="hex", 
    default=False)

parser.add_option("-q", 
    help="Quiet mode", 
    action="store_true", 
    dest="quiet", 
    default=False)

parser.add_option("-e", 
    metavar="FLOAT", 
    help="Number of bits of entropy per byte (default is %default)", 
    action="store", 
    type="float", 
    dest="entropy", 
    default=5.0)

#parser.add_option("-r",
#    metavar="INT", 
#    help="Serial port's bitrate (default is %default)", 
#    action="store", 
#    type="int", 
#    dest="rate", 
#    default=57600)
#
#parser.add_option("-d",  
#    metavar="DEVICE", 
#    help="Serial port device (default is %default)", 
#    action="store", 
#    type="string", 
#    dest="device", 
#    default="/dev/ttyUSB0")

(options, args) = parser.parse_args()

entropy_per_byte = options.entropy
hash_len = sha.digest_size
read_len = int((hash_len * 8) / entropy_per_byte)

if options.data_length == None:
  parser.exit(1, "Error: You must specify a data length\n")

data_length = options.data_length
if data_length[-1] == 'b':
  data_length = int(data_length[:-1])
  if data_length % 8 != 0:
    parser.exit(1, "Error: data length must be a multiple of 8 bits\n")
  else:
    data_length /= 8
elif data_length[-1] == 'K':
  data_length = int(data_length[:-1])*1024
elif data_length[-1] == 'M':
  data_length = int(data_length[:-1])*1024*1024
else:
  try:
    data_length = int(data_length)
  except ValueError:
    parser.exit(1, "Error: '%s' is not a valid data length\n" % data_length)

output_count = 0
output_file = sys.stdout
if options.output_file:
  output_file = open(options.output_file, "w")

try:
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  sock.connect(("localhost", 9876))
  sock_file = sock.makefile("rb") # buffered
except socket.error:
  parser.exit(1, "Error: Failed to connect to button_box_server\n")


start_time = time.time()
next_check_time = start_time+1

cont = True
while cont and (data_length == None or output_count < data_length):
  block = ""
  while len(block) < read_len:
    try:
      sock.send("Q")
      # expect to receive 'Qx\n' where x is a random byte
      buf = sock_file.readline()
      while not buf[0] == 'Q':
        buf = sock_file.readline()
      block += buf[1] # get the random byte
    except KeyboardInterrupt:
      cont = False
      break # exit the inner loop
  if not cont:
    break # exit the outer loop

  to_write = hash_len
  if output_count + hash_len > data_length:
    to_write = data_length - output_count
  if options.hex:
    hash = sha.new(block).hexdigest()
    output_file.write(hash[:to_write*2])
  else:
    hash = sha.new(block).digest()
    output_file.write(hash[:to_write])
  output_file.flush()

  output_count += to_write
  if not options.quiet and time.time() > next_check_time and output_count > 0:
    next_check_time = time.time()+1
    elapsed = time.time() - start_time
    eta = elapsed / (float(output_count) / data_length) - elapsed
    sys.stderr.write("\r%9d bytes (%3d%%), ETA: %02d:%02d (%.3f kbps)" % (output_count, 100.0 * output_count/data_length, eta/60, eta%60, 8*output_count/(1000.0*elapsed)))
    sys.stderr.flush()

if options.hex:
  output_file.write("\n")
if not options.quiet:
  elapsed_time = time.time() - start_time
  sys.stderr.write("\nElapsed time: %02d:%02d, rate: %.3f kbps, %d bytes written\n" %
      (elapsed_time/60, elapsed_time % 60, 8*output_count/(1000.0*elapsed_time), output_count))