#! /usr/bin/python2
# -*- coding: utf-8 -*-
import sys
import os.path
import subprocess
import time
import re
import shlex
import datetime
import shutil

from optparse import OptionParser

eid_list       = []
exports_list   = []
client_list    = []
count_list   = []
max_read_list  = []
max_write_list = []
active_list    = []

#___________________________
def append_target(eid,exports):

  # Check this eid is not yet referenced
  for other_eid in eid_list:
    if eid == other_eid: return
      
  eid_list.append(eid)
  exports_list.append(exports)
  
  active_export = get_active_export(exports)
  active_list.append(active_export)
  
  clients,count  = get_client_list(active_export,eid)
  client_list.append(clients)
  count_list.append(count)

  print "  Clients of eid %s : "%(eid)
  for client in clients.split("-i"):
    words = client.split()
    if len(words) >= 3:
      print "  - %s:%s"%(words[0],words[2])
  print
  
  max_read_list.append(0)
  max_write_list.append(0)
#___________________________
# converts string <inBytes> to integer 
# in the unit given by <to>.
# <inBytes>  is 
# [<spaces>]<numeric>[{.|,}<numeric>][<spaces>]<unit>
# <unit> is {B|K|KB|M|MB|G|GB|T|TB|P|PB}
# <to> is <unit>
def common_convert_Bytes(inBytes,to="B"):
  inBytes = inBytes.replace(" ", "")

  if "." in inBytes:
    before = int(inBytes.split('.')[0])
    after  = inBytes.split('.')[1]
    unit   = re.split('(\d+)',after)[2]
    after  = re.split('(\d+)',after)[1]
  elif "," in inBytes:
    before = int(inBytes.split(',')[0])
    after  = inBytes.split(',')[1]
    unit   = re.split('(\d+)',after)[2]
    after  = re.split('(\d+)',after)[1]
  else:
    unit   = re.split('(\d+)',inBytes)[2]
    before = re.split('(\d+)',inBytes)[1]
    after = "0"

  if   len(after) == 1: after=int(after)*100
  elif len(after) == 2: after=int(after)*10
  elif len(after) == 3: after=int(after)

  if   unit == "KB" or unit == "K": val = (int(before)*1000 + after)
  elif unit == "MB" or unit == "M": val = (int(before)*1000 + after) * 1000
  elif unit == "GB" or unit == "G": val = (int(before)*1000 + after) * 1000000
  elif unit == "TB" or unit == "T": val = (int(before)*1000 + after) * 1000000000
  elif unit == "PB" or unit == "P": val = (int(before)*1000 + after) * 1000000000000
  else: val = int(before)
  
  if to == "K" or to =="KB": return val/1000
  if to == "M" or to =="MB": return val/1000000  
  if to == "G" or to =="GB": return val/1000000000
  if to == "T" or to =="TB": return val/1000000000000
  if to == "P" or to =="PB": return val/1000000000000000
  return val
#______________________  
def common_bytes_with_unit(inBytes):
  if inBytes < 0: sign = "-"
  else:           sign = ""
  
  string="%s"%(inBytes)
  value = common_convert_Bytes(string,"B")

  if value < 1000: return "%s%4d  "%(sign,value)
  
  if value < 1000000:
    if value >= (100*1000): return "%s%4d K"%(sign,value/1000)
    if value >=  (10*1000): return "%s%2d.%1d K"%(sign,value/1000,(value % 1000) / 100)
    return "%s%1d.%2.2d K"%(sign,value/1000,(value % 1000) / 10)   
  
  if value < 1000000000:  
    if value >= (100*1000000): return "%s%4d M"%(sign,value/1000000)
    if value >=  (10*1000000): return "%s%2d.%1d M"%(sign,value/1000000,(value % 1000000) / 100000)
    return "%s%1d.%2.2d M"%(sign,value/1000000,(value % 1000000) / 10000)   
  
  if value < 1000000000000:  
    if value >= (100*1000000000): return "%s%4d G"%(sign,value/1000000000)
    if value >=  (10*1000000000): return "%s%2d.%1d G"%(sign,value/1000000000,(value % 1000000000) / 100000000)
    return "%s%1d.%2.2d G"%(sign,value/1000000000,(value % 1000000000) / 10000000)   
  
  if value < 1000000000000000:  
    if value >= (100*1000000000000): return "%s%4d T"%(sign,value/1000000000000)
    if value >=  (10*1000000000000): return "%s%2d.%1d T"%(sign,value/1000000000000,(value % 1000000000000) / 100000000000)
    return "%s%1d.%2.2d T"%(sign,value/1000000000000,(value % 1000000000000) / 10000000000)   
  
  if value < 1000000000000000000:  
    if value >= (100*1000000000000000): return "%s%4d P"%(sign,value/1000000000000000)
    if value >=  (10*1000000000000000): return "%s%2d.%1d P"%(sign,value/1000000000000000,(value % 1000000000000000) / 100000000000000)
  
  return "%s%1s.%2.2d P"%(sign,value/1000000000000000,(value % 1000000000000000) / 10000000000000)   

#______________________________
def is_this_export_active(export):

  string="rozodiag -i %s -T export:1 -c up"%(export)
  parsed = shlex.split(string)
  cmd = subprocess.Popen(parsed, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  for line in cmd.stdout:
      if "uptime" in line:
        return True
  return False
#______________________________
def get_active_export(exports,active_export=None):

  #
  # Last active export is given
  # Check it 1rst
  #
  if active_export != None:
    if is_this_export_active(active_export) == True: return active_export
  
  #
  # No last active given
  # or last active no more active (switchover
  #  
  for export in exports.split('/'):

    #
    # The last active is no more active
    #
    if export == active_export: continue
    
    if is_this_export_active(export) == True: 
      return export
      
  return active_export
#______________________________
def get_eid_and_exports_from_path(path):
  eid=None
  string="attr -qg rozofs.export %s"%(path)
  parsed = shlex.split(string)
  cmd = subprocess.Popen(parsed, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  for line in cmd.stdout:
    exports=line.split()[0]
    eid=line.split()[1]
    return eid,exports
  syntax( "%s is not a RozoFS mountpoint"%(path))
            
#______________________________
def check_netmwork_filters(address):

  try:
    ip = get_ip(address)
  except: 
    return False

  if len(NipLowArray) != 0:
    idx = 0
    for low in NipLowArray:
      high = NipHighArray[idx]
      if ip >= low and ip <= high: 
        if options.detail: print RED+"%s rejected %s"%(address,NipRule[idx])+NORMAL 
        return False
      idx += 1
          
  if len(ipLowArray) != 0:
    idx = 0
    for low in ipLowArray:
      high = ipHighArray[idx]
      if ip >= low and ip <= high: 
        if options.detail: print GREEN+"%s accepted %s"%(address,ipRule[idx])+NORMAL     
        return True
      idx += 1
  else:
    if options.detail: print GREEN+"%s accepted"%(address)+NORMAL 
    return True
    
  if options.detail: print RED+"%s rejected"%(address)+NORMAL 
  return False
  
#______________________________
def get_client_list(active_export,eid):
  
  count = int(0) 
  idx = int(eid) % 8
  clients=""
  string="rozodiag -i %s -T export:%s -c client"%(active_export,idx)
  parsed = shlex.split(string)
  cmd = subprocess.Popen(parsed, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  for line in cmd.stdout:
    words = line.split('|')
    try:
      int(words[3])
      delay = int(words[2])
      if delay > int(300) : continue
    except:
      continue  
    ip = words[4].split(':')[0]
    if check_netmwork_filters(ip) == False: continue
    port = words[4].split(':')[1]
    string="rozodiag -i %s -p %s -c mount"%(ip,port)
    parsed = shlex.split(string)
    cmd = subprocess.Popen(parsed, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    for line in cmd.stdout:
      if line.split()[0] != "eid": continue
      if line.split()[2] != eid: continue
      clients = clients + " -i %s -p %s"%(ip,port)  
      count = count + int(1)   
  return clients,count      
#______________________________
def get_ip(ipString):
 if len(ipString.split('.')) != 4: int(a)
 ip = int(0)
 for val in ipString.split('.'):
   ip = ip * int(256) + int(val)
 return ip  

#______________________________________    
def get_line_index(val):
  if val == " Avg ": return 10
  try:
    return -int(val) - 1
  except:
    return -1  
#______________________________________    
def get_color_code(val):
  if val == "red": return RED
  if val == "green": return GREEN
  if val == "yellow": return YELLOW
  if val == "blue": return BLUE
  if val == "purple": return PURPLE
  if val == "cyan": return CYAN
  if val == "orange": return ORANGE
  return WHITE    
#______________________________________    
def syntax(string=None):
  if string != None:
    print "\n" + RED + BOLD + string + " !!!" + NORMAL + "\n"
    
  print "Usage: " + BOLD + "rozo_throughput "+ YELLOW +"[OPTIONS]" + CYAN + " [FILTER] " + GREEN + "[MOUNTPOINTS]"+ NORMAL
  print
  print YELLOW +"  OPTIONS = {--period <seconds> | --readcolor <color> | --writecolor <color> | --history}"+ NORMAL
  print YELLOW +"    -p,--period <seconds>"+ NORMAL +"   enables to set the refresh period to <seconds> seconds."
  print YELLOW +"    -r,--readcolor <color>"+ NORMAL +"  to change read color."
  print YELLOW +"    -w,--writecolor <color>"+ NORMAL +" to change write color."
  print YELLOW +"    -H,--history"+ NORMAL +"            displays a 60 minutes and a 60 hours history."
  print
  print CYAN + "  FILTER = {--ipRange <IP1@-IP2@> ... | --NipRange <IP1@-IP2@> ...}"+ NORMAL
  print CYAN + "    -i,--ipRange  <IP1@-IP2@>"+ NORMAL  +"  excludes clients whose IP address is out of IP1-IP2 range." 
  print CYAN + "    -I,--NipRange <IP1@-IP2@>"+ NORMAL  +"  excludes clients whose IP address is within IP1-IP2 range." 
  print
  print GREEN + "  MOUNTPOINTS = [<mnt>] ..."+ NORMAL
  print GREEN + "    mnt = "+ NORMAL  +"RozoFS mount point"
  print "  When no mount point is provided, current path is taken as input mount point."
  print 
  print "examples:"
  print "  rozo_throughput /mnt/rfs1" 
  print "  rozo_throughput -p 10 /mnt/rfs /mnt/private" 
  print 
  sys.exit(-1)
#______________________________________    

RED='\033[91m\033[40m'
GREEN='\033[92m\033[40m'
YELLOW='\033[93m\033[40m'
BLUE='\033[94m\033[40m'
PURPLE='\033[95m\033[40m'
CYAN='\033[96m\033[40m'
WHITE='\033[97m\033[40m'
ORANGE='\033[33m\033[40m'

BOLD='\033[1m'
REVERSE='\033[7m'
NORMAL='\033[0m'

color_list=["red","green","yellow","blue","purple","cyan","white","orange"]
time_unit_list=["s","m","h"]

parser = OptionParser()
parser.add_option("-p","--period", action="store",type="string", dest="period", help="Periodicity of the command")
parser.add_option("-d","--detail", action="store_true",default=False, dest="detail", help="Detailed per client throughput display.")
parser.add_option("-o","--output", action="store",type="string", dest="output", help="Output file")
parser.add_option("-r","--readcolor", type = "choice", action="store",dest ="rcol", choices = color_list, default = "cyan", help = "Read color within %s."% (color_list))
parser.add_option("-w","--writecolor", type = "choice", action="store",dest ="wcol", choices = color_list, default = "yellow", help = "Write color within %s."% (color_list))
parser.add_option("-H","--history", action="store_true",default=False, dest="history", help="Display a 60 minutes and 60 hours history.")
parser.add_option("-i","--ipRange", action="append",type="string", dest="iprange", help="Allowed IP range (one IP address or a range like 192.168.1.21-192.168.1.28)")
parser.add_option("-I","--NipRange", action="append",type="string", dest="Niprange", help="Rejected IP range (one IP address or a range like 192.168.10.30-192.168.10.50)")

(options, args) = parser.parse_args()
period  = int(30)
if options.period != None:
  try:
    period = int(options.period)
  except:
    syntax( "Bad period value %s"%(options.period))


# 
# Choosen colors
# 
READCOLOR  = get_color_code(options.rcol) 
WRITECOLOR = get_color_code(options.wcol) 
    
ipLowArray  = []
ipHighArray = []
ipRule      = []
NipLowArray  = []
NipHighArray = []
NipRule      = []

if options.iprange != None:
  for iprange in options.iprange:
    if len(iprange.split('-')) < 2:
      ipStringLow  = iprange
      ipStingHigh  = iprange
    else:      
      ipStringLow  = iprange.split('-')[0]
      ipStingHigh  = iprange.split('-')[1]
    try:
      ipLow  = get_ip(ipStringLow)
      ipHigh = get_ip(ipStingHigh)
    except:
      syntax( "Bad IP range %s"%(iprange))
    if ipLow > ipHigh:
      syntax( "Low IP is greater than high IP in range %s"%(iprange))
    ipLowArray.append(ipLow)
    ipHighArray.append(ipHigh)
    ipRule.append(iprange)

if options.Niprange != None:
  for iprange in options.Niprange:
    if len(iprange.split('-')) < 2:
      ipStringLow  = iprange
      ipStingHigh  = iprange
    else:
      ipStringLow  = iprange.split('-')[0]
      ipStingHigh  = iprange.split('-')[1]
    try:
      ipLow  = get_ip(ipStringLow)
      ipHigh = get_ip(ipStingHigh)
    except:
      syntax( "Bad IP range %s"%(iprange))
    if ipLow > ipHigh:
      syntax( "Low IP is greater than high IP in range %s"%(iprange))
    NipLowArray.append(ipLow)
    NipHighArray.append(ipHigh)      
    NipRule.append(iprange)

#
# input RozoFS mount point is either in args or current path
#
print
if len(args) == 0:
  eid,exports = get_eid_and_exports_from_path(os.getcwd())
  append_target(eid,exports)
else:  
  for arg in args:
    eid,exports = get_eid_and_exports_from_path(arg)
    append_target(eid,exports)

#
# Redirect output
#  
if options.output != None:
  sys.stdout = open(options.output, 'w')
  if os.fork(): sys.exit()
  
#
# Change current directory 
# Not to block the file system
#
os.chdir("/")  
history = ""

loop = int(0)
while True :   

  #
  # Every 10 loops, rebuild the list of clients
  #
  loop = loop + int(1)
  if loop == 12: loop = 0
      
  for index in range(len(eid_list)):

    eid           = eid_list[index]
    exports       = exports_list[index]
    active_export = active_list[index]
    clients       = client_list[index]
    count         = count_list[index]
    max_read      = max_read_list[index]
    max_write     = max_write_list[index]    

    if loop == 11:
      active_export = get_active_export(exports,active_export)
      clients,count = get_client_list(active_export,eid)
      
      count_list[index]   = count  
      active_list[index]  = active_export 
      client_list[index]  = clients


    read  = [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
    write = [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
    read_avg   = [0,0,0,0,0,0]  
    write_avg  = [0,0,0,0,0,0]  
    local_max_read  = 0
    local_max_write = 0
    local_max_read_avg  = 0
    local_max_write_avg = 0
    is_read_max  = False
    is_write_max = False

    #
    # Get current time
    #
    zedate = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    #
    # Run the rozodiag command against all cliets
    #
    string="rozodiag %s -c throughput col 6 avg "%(clients)
    if options.history == True: 
      if history == "": 
        history = " m"
      else:      
        history = " h"
    string += history
    parsed = shlex.split(string)
    cmd = subprocess.Popen(parsed, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    #
    # Parse output and cumulate the throughputs
    #
    for line in cmd.stdout:
      line = line.split('\n')[0]
      if options.detail == True: print "%s"%(line)
      line = line.split('|')
      if len(line) < 23: continue 
      idx = get_line_index(line[1])  
      if idx == -1: continue

      #
      # Average line
      #
      if idx == 10:

        read_avg[0]   = read_avg[0]  + common_convert_Bytes(line[2])
        write_avg[0]  = write_avg[0]  + common_convert_Bytes(line[3])
        if read_avg[0] > local_max_read_avg: local_max_read_avg = read_avg[0] 
        if write_avg[0] > local_max_write_avg: local_max_write_avg = write_avg[0] 

        read_avg[1]   = read_avg[1]  + common_convert_Bytes(line[6])
        write_avg[1]  = write_avg[1]  + common_convert_Bytes(line[7])
        if read_avg[1] > local_max_read_avg: local_max_read_avg = read_avg[1] 
        if write_avg[1] > local_max_write_avg: local_max_write_avg = write_avg[1] 

        read_avg[2]   = read_avg[2]  + common_convert_Bytes(line[10])
        write_avg[2]  = write_avg[2] + common_convert_Bytes(line[11])
        if read_avg[2] > local_max_read_avg: local_max_read_avg = read_avg[2] 
        if write_avg[2] > local_max_write_avg: local_max_write_avg = write_avg[2] 

        read_avg[3]   = read_avg[3]  + common_convert_Bytes(line[14])
        write_avg[3]  = write_avg[3] + common_convert_Bytes(line[15])
        if read_avg[3] > local_max_read_avg: local_max_read_avg = read_avg[3] 
        if write_avg[3] > local_max_write_avg: local_max_write_avg = write_avg[3] 

        read_avg[4]   = read_avg[4]   + common_convert_Bytes(line[18])
        write_avg[4]  = write_avg[4]  + common_convert_Bytes(line[19])
        if read_avg[4] > local_max_read_avg: local_max_read_avg = read_avg[4] 
        if write_avg[4] > local_max_write_avg: local_max_write_avg = write_avg[4] 

        read_avg[5]   = read_avg[5]  + common_convert_Bytes(line[22])
        write_avg[5]  = write_avg[5] + common_convert_Bytes(line[23])
        if read_avg[5] > local_max_read_avg: local_max_read_avg = read_avg[5] 
        if write_avg[5] > local_max_write_avg: local_max_write_avg = write_avg[5] 
        continue

      #
      # Normal line
      #  
      read[idx]     = read[idx]      + common_convert_Bytes(line[2])
      write[idx]    = write[idx]     + common_convert_Bytes(line[3])
      read[idx+10]  = read[idx+10]   + common_convert_Bytes(line[6])
      write[idx+10] = write[idx+10]  + common_convert_Bytes(line[7])
      read[idx+20]  = read[idx+20]   + common_convert_Bytes(line[10])
      write[idx+20] = write[idx+20]  + common_convert_Bytes(line[11])
      read[idx+30]  = read[idx+30]   + common_convert_Bytes(line[14])
      write[idx+30] = write[idx+30]  + common_convert_Bytes(line[15])
      read[idx+40]  = read[idx+40]   + common_convert_Bytes(line[18])
      write[idx+40] = write[idx+40]  + common_convert_Bytes(line[19])
      read[idx+50]  = read[idx+50]   + common_convert_Bytes(line[22])
      write[idx+50] = write[idx+50]  + common_convert_Bytes(line[23])

    #
    # Search for the maximum values
    #
    for idx in range(60):
      if read[idx] > local_max_read: local_max_read = read[idx]
      if write[idx] > local_max_write: local_max_write = write[idx]    
    if local_max_write > max_write:  
      max_write = local_max_write
      is_write_max = True
    if local_max_read  > max_read:   
      max_read  = local_max_read    
      is_read_max = True

    #
    # Prepare display header
    #
    LINE = zedate + " - %s RozoFS clients on "%(count) + GREEN +" eid " + eid + NORMAL  
    if history == " m":
      LINE = LINE + " - Last " + REVERSE + "60 minutes" + NORMAL + " history."  
    elif history == " h": 
      LINE = LINE + " - Last " + REVERSE + "60 hours" + NORMAL + " history."
    else: 
      LINE = LINE + " - max " 
      LINE += READCOLOR 
      if is_read_max == True: LINE += REVERSE
      LINE += common_bytes_with_unit(max_read) + NORMAL + '/' +  WRITECOLOR  
      if is_write_max == True: LINE += REVERSE
      LINE += common_bytes_with_unit(max_write) + NORMAL
    print LINE  

    print " __Read__ __Write_  __Read__ __Write_  __Read__ __Write_  __Read__ __Write_  __Read__ __Write_  __Read__ __Write_ "

    for idx in range(10):
      LINE=""
      for cols in range(6):
        IDX = idx + cols*10
        LINE += "|"
        val = READCOLOR + " %6s "%(common_bytes_with_unit(read[IDX]))  
        if read[IDX] == 0:                     LINE += val  
        elif read[IDX] == local_max_read:      LINE += REVERSE + val  
        elif read[IDX] >= local_max_read_avg:  LINE += BOLD + val  
        else:                                  LINE += val  
        LINE += NORMAL + "|"
        val = WRITECOLOR + " %6s "%(common_bytes_with_unit (write[IDX]))      
        if write[IDX] == 0:                     LINE += val  
        elif write[IDX] == local_max_write:     LINE += REVERSE + val 
        elif write[IDX] >= local_max_write_avg: LINE += BOLD + val 
        else:                                   LINE += val 
        LINE += NORMAL + "|"
      print LINE    
    print "|________|________||________|________||________|________||________|________||________|________||________|________|"

    LINE=""

    for IDX in range(6):
      LINE += "|"
      val = READCOLOR + " %6s "%(common_bytes_with_unit(read_avg[IDX])) 
      if read_avg[IDX] == 0:                    LINE += val    
      elif read_avg[IDX] == local_max_read_avg: LINE += REVERSE + val  
      else:                                     LINE += val   
      LINE += NORMAL + "|"
      val = WRITECOLOR + " %6s "%(common_bytes_with_unit(write_avg[IDX]))     
      if write_avg[IDX] == 0:                     LINE += val    
      elif write_avg[IDX] == local_max_write_avg: LINE += REVERSE + val  
      else:                                       LINE += val    
      LINE += NORMAL + "|"
    print LINE    

    sys.stdout.flush()
    
    max_read_list[index]  = max_read
    max_write_list[index] = max_write  
    
  if history == " h": break;
  if history != " m": 
    if period == 0: break
    time.sleep(period)
    
