#!/usr/bin/python
import getopt
import sys
import re

#--- sub
def usage():
  print 'This program reads a table of VO data, and uses it to'
  print 'construct some products.'
  print ''
  print 'The options to this program are:'
  print ' -c   --conf         <file>   conf file to read  (def: vo_data.conf)'
  print ' -u   --users        <file>   users conf file    (def: users.conf.gen)'
  print ' -n   --condor       <file>   condor conf file   (def: condor.conf.gen)'
  print ' -r   --groupmapfile <file>   groupmapfile       (def: groupmapfile.gen)'
  print ' -m   --gridmapfile  <file>   gridmapfile        (def: grid-mapfile.gen)'
  print ' -a   --arguswn      <file>   arguswn conf       (def: arguswn.conf.gen)'
  print ''

#--- sub
def initOptions(o):

  """ Sets the defaults and read the options """

  # Set defaults
  o['vo_data_input']      = 'vo_data.conf'
  o['users_conf_output']   = 'users.conf.gen'
  o['condor_conf_output']  = 'condor.conf.gen'
  o['groupmap_output']     = 'groupmapfile.gen'
  o['gridmap_output']       = 'grid-mapfile.gen'
  o['arguswn_conf_output'] = 'arguswn.conf.gen'
 
  # Read the options 
  try:
    options, remainder = getopt.getopt(sys.argv[1:], 'c:u:n:r:m:a:', [
      'conf=', 'users=', 'condor=', 'groupmapfile=','gridmapfile=','arguswn=' ])

  except getopt.GetoptError:
    usage(); sys.exit(1)

  if len(options) <= 0:
    usage(); sys.exit(1)
 
  # Store the options 
  for opt, arg in options:
    if opt in ('-c', '--conf'):
      o['vo_data_input'] =  arg
    elif opt in ('-u', '--users'):
      o['users_conf_output'] = arg
    elif opt in ('-n', '--condor'):
      o['condor_conf_output'] = arg
    elif opt in ('-r', '--groupmapfile'):
      o['groupmap_output'] = arg
    elif opt in ('-m', '--gridmapfile'):
      o['gridmap_output'] = arg
    elif opt in ('-a', '--arguswn'):
      o['arguswn_conf_output'] = arg
  
#--- class
class UsersConf:

  """ Produces a users.conf file """

  # Write the users.conf
  def write_users (self,vo_list,file,start_uid,start_add_uid, start_sguid):

    try:
      f = open(file, 'w')
    except IOError:
      print('Cannot open users.conf file '); sys.exit(1)

    # vouid - first UID in VO's range
    # voadduid - first UID in the additional range for VO
    vouid = start_uid
    voadduid = start_add_uid
    vosguid = start_sguid

    # For over the VOs
    for vo in vo_list:

      # Write block STD USERS
      nuid = vouid
      for jj in range(vo.users):
        record =  "%d:%s%03d:%d:%s:%s::\n" % (nuid, vo.name5, jj + 1, vo.gid,vo.name8,vo.name)
        match_obj = re.match('^DUMMY',vo.blockname)
        if match_obj is None:
          f.write(record)

        nuid = nuid + 1
    
      # Write block of PRODUCTION USERS
      nuid = vouid + 300
      n=vo.prd
      for jj in range(n ):
        record =  "%d:%s%02d:%d,%d:%s,%s:%s:%s:\n" % (nuid, 'prd' + vo.name3 , jj + 1, vo.pgid, vo.gid, vo.name5 + 'prd',vo.name8,vo.name,'prd')
        match_obj = re.match('^DUMMY',vo.blockname)
        if match_obj is None:
          f.write(record)
        nuid = nuid + 1

      # Write block of SGM USERS
      nuid = vouid + 400
      n=vo.sgm
      for jj in range(n ):
        record =  "%d:%s%02d:%d,%d:%s,%s:%s:%s:\n" % (nuid,'sgm' + vo.name3, jj + 1,vo.sgid,vo.gid,vo.name5 + 'sgm',vo.name8,vo.name,'sgm')
        match_obj = re.match('^DUMMY',vo.blockname)
        if match_obj is None:
          f.write(record)
        nuid = nuid + 1

      # Maybe write block of PILOT USERS (in additional range). Only happens if vo.pl > 0
      nuid = voadduid
      n = vo.pil
      for jj in range(n ):
        record =  "%d:%s%02d:%d,%d:%s,%s:%s:%s:\n" % (nuid,'pil' + vo.name3, jj + 1,vo.pilgid,vo.gid,vo.name5 + 'pil',vo.name8,vo.name,'pilot')
        match_obj = re.match('^DUMMY',vo.blockname)
        if match_obj is None:
          f.write(record)
        nuid = nuid + 1

      # Maybe write a block of SubGroup users. Only happens if sub groups defined for this VO
      sg_list = vo.sg_list
      block_nuid = vosguid
      for sg in sg_list:
        nuid = block_nuid
        n = sg.users
        for jj in range(n):
          record =  "%d:%s%02d:%d,%d:%s,%s:%s:%s:\n" % (nuid,sg.name, jj + 1,  sg.guid,vo.gid,sg.name + 'sg',vo.name8,vo.name,'subgroup')
          match_obj = re.match('^DUMMY',vo.blockname)
          if match_obj is None:
            f.write(record)
          nuid = nuid + 1
        block_nuid = block_nuid + 100
    
      # set vouid and voadduid for next VO
      vouid = vouid + 500
      voadduid = voadduid + 250
      vosguid = vosguid + 250

    f.close()

#--- class
class CondorConf:

  """ Produces a condor.conf file """

  # Write the condor.conf
  def write_condor (self,vo_list,file):

    try:
      f = open(file, 'w')
    except IOError:
      print('Cannot open condor.conf file '); sys.exit(1)

    f.write('# Group names\n')
    f.write('GROUP_NAMES = \\\n')
    # 
    record =  "\tgroup_HIGHPRIO" 
    f.write(record + ',  \\\n')

    # For over the VOs
    for vo in vo_list:

      voName = re.sub(r"\.", "_", vo.name.upper())

      # Write generic group record
      record =  "\tgroup_%s" % (voName)
      match_obj = re.match('^DUMMY',vo.blockname)
      if match_obj is None:
        f.write(record + ',  \\\n')

    f.write('\n')
    f.close()

    try:
      f = open(file, 'a')
    except IOError:
      print('Cannot open condor.conf file '); sys.exit(1)

    f.write('# Fairshares\n')

    record =  "GROUP_QUOTA_DYNAMIC_group_HIGHPRIO  = 0.05\n"
    f.write(record + '\n')

    # For over the VOs
    for vo in vo_list:

      voName = re.sub(r"\.", "_", vo.name.upper())

      # Write generic group record
      record =  "GROUP_QUOTA_DYNAMIC_group_%s = %5.2f" % (voName, float(vo.share))
      match_obj = re.match('^DUMMY',vo.blockname)
      if match_obj is None:
        f.write(record + '\n')

    f.write('\n')
    f.close()
    #---------------
    try:
      f = open(file, 'a')
    except IOError:
      print('Cannot open condor.conf file '); sys.exit(1)

    f.write('# Priority factors\n')

    record =  "DEFAULT_PRIO_FACTOR = 100000.00"
    f.write(record + '\n')
    record =  "GROUP_PRIO_FACTOR_group_HIGHPRIO  = 1.00\n"
    f.write(record + '\n')

    # For over the VOs
    for vo in vo_list:

      voName = re.sub(r"\.", "_", vo.name.upper())

      # Write generic group record
      record =  "GROUP_PRIO_FACTOR_group_%s = %5.2f" % (voName, 10000)
      match_obj = re.match('^DUMMY',vo.blockname)
      if match_obj is None:
        f.write(record + '\n')

    f.write('\n')
    f.close()

#--- class
class GroupMapFile:

  """ Produces a groupmapfile """
  def __init__(self):
    x = 1

  # Write groupmapfile
  def write_groupmapfile(self,vo_list,file):

  # Produce groupmapfile format


    try:
      f = open(file, 'w')
    except IOError:
      print('Cannot open groupmapfile '); sys.exit(1)

    # Go over all the VOs
    for vo in vo_list:

      # Omit DUMMY...
      match_obj = re.match('^DUMMY',vo.blockname) 
      if match_obj is None:

        # Write this VO

        f.write('"/' + vo.name + '/sgm/Role=NULL/Capability=NULL" ' + vo.name5 + 'sgm\n')
        f.write('"/' + vo.name + '/sgm" ' + vo.name5 + 'sgm\n')
  
        f.write('"/' + vo.name + '/lcgprod/Role=NULL/Capability=NULL" ' + vo.name5 + 'prd\n')
        f.write('"/' + vo.name + '/lcgprod" ' + vo.name5 + 'prd\n')

        f.write('"/' + vo.name + '/Role=lcgadmin/Capability=NULL" ' + vo.name5 + 'sgm\n')
        f.write('"/' + vo.name + '/Role=lcgadmin" ' + vo.name5 + 'sgm\n')

        f.write('"/' + vo.name + '/Role=production/Capability=NULL" ' + vo.name5 + 'prd\n')
        f.write('"/' + vo.name + '/Role=production" ' + vo.name5 + 'prd\n')

        # Maybe write pilots
        if re.match('yes',vo.pilot):
          f.write('"/' + vo.name + '/Role=pilot/Capability=NULL" ' + vo.name5 + 'pil\n')
          f.write('"/' + vo.name + '/Role=pilot" ' + vo.name5 + 'pil\n')

        # Maybe write a block of SubGroup users. Only happens if sub groups defined for this VO
        sg_list = vo.sg_list
        for sg in sg_list:
          f.write('"' + sg.fqdn + '" ' + sg.name + 'sg\n')

        f.write('"/' + vo.name + '/Role=NULL/Capability=NULL" ' + vo.name8 + '\n')
        f.write('"/' + vo.name + '" ' + vo.name8 + '\n')
        f.write('"/' + vo.name + '/*/Role=NULL/Capability=NULL" ' + vo.name8 + '\n')
        f.write('"/' + vo.name + '/*" ' + vo.name8 + '\n')

    f.close()

#--- class
class GridMapFile:

  """ Produces a groupmapfile """
  def __init__(self):
    x = 1

  # Write gridmapfile
  def write_gridmapfile(self,vo_list,file):

  # 

    try:
      f = open(file, 'w')
    except IOError:
      print('Cannot open gridmapfile '); sys.exit(1)

    # Go over all the VOs
    for vo in vo_list:

      # Omit DUMMY...
      match_obj = re.match('^DUMMY',vo.blockname) 
      if match_obj is None:

        # Write this VO

        f.write('"/' + vo.name + '/sgm/Role=NULL/Capability=NULL" .sgm' + vo.name3 + '\n')
        f.write('"/' + vo.name + '/sgm" .sgm' + vo.name3 + '\n')
  
        f.write('"/' + vo.name + '/lcgprod/Role=NULL/Capability=NULL" .prd' + vo.name3 + '\n')
        f.write('"/' + vo.name + '/lcgprod" .prd' + vo.name3 + '\n')

        f.write('"/' + vo.name + '/Role=lcgadmin/Capability=NULL" .sgm' + vo.name3 + '\n')
        f.write('"/' + vo.name + '/Role=lcgadmin" .sgm' + vo.name3 + '\n')

        f.write('"/' + vo.name + '/Role=production/Capability=NULL" .prd' + vo.name3 + '\n')
        f.write('"/' + vo.name + '/Role=production" .prd' + vo.name3 + '\n')

        # Maybe write pilots
        if re.match('yes',vo.pilot):
          f.write('"/' + vo.name + '/Role=pilot/Capability=NULL" .pil' + vo.name3 + '\n')
          f.write('"/' + vo.name + '/Role=pilot" .pil' + vo.name3 + '\n')

        # Maybe write a block of SubGroup users. Only happens if sub groups defined for this VO
        sg_list = vo.sg_list
        for sg in sg_list:
          f.write('"' + sg.fqdn + '" .' + sg.name + '\n')

        f.write('"/' + vo.name + '/Role=NULL/Capability=NULL" .' + vo.name5 + '\n')
        f.write('"/' + vo.name + '" .' + vo.name5 + '\n')
        f.write('"/' + vo.name + '/*/Role=NULL/Capability=NULL" .' + vo.name5 + '\n')
        f.write('"/' + vo.name + '/*" .' + vo.name5 + '\n')

    f.close()


#--- class
class ArguswnConf:

  """ Produces a section of argus policy for the WNs """
  def __init__(self):
    x = 1

  # Write argus policy file
  def write_arguswn(self,vo_list,file):

    try:
      f = open(file, 'w')
    except IOError:
      print('Cannot open arguswn.conf file '); sys.exit(1)

    # Go over all the VOs
    for vo in vo_list:

      # Omit DUMMY...
      match_obj = re.match('^DUMMY',vo.blockname) 
      if match_obj is None:

        # Write this VO

        f.write('    rule permit {pfqan = "/' + vo.name + '\" }\n') 
        f.write('    rule permit {pfqan = "/' + vo.name + '/Role=lcgadmin\" }\n') 
        f.write('    rule permit {pfqan = "/' + vo.name + '/Role=production\" }\n') 

        # Maybe write pilots
        if re.match('yes',vo.pilot):
          f.write('    rule permit {pfqan = "/' + vo.name + '/Role=pilot\" }\n') 
        sg_list = vo.sg_list;
        for sg in sg_list:
          f.write('    rule permit {pfqan = "' + sg.fqdn + '\" }\n') 

    f.close()


#--- class
class SubGroup:
  
  """ Holds a SubGroup """

  # Constructor, with initial attribute values
  def __init__(self,attribs):
    self.name = attribs[0]
    self.users = int(attribs[1])
    self.fqdn = attribs[2]
    self.share = attribs[3]
    self.guid = int(attribs[4])

#--- class
class Vo:

  """ Holds a VO """

  # Constructor, with initial attribute values
  def __init__(self,fields):

    # Names, gids etc.

    self.sg_list = []

    self.blockname = fields[0]
    self.gid = int(fields[1])
    self.pgid = int(fields[2])
    self.sgid = int(fields[3])
    self.pilgid = int(fields[4])
    self.name = fields[5]
    self.name3 = fields[6]
    self.name5 = fields[7]
    self.name8 = fields[8]
    self.users = int(fields[9])
    self.prd = int(fields[10])
    self.sgm = int(fields[11])
    self.pil = int(fields[12])
    self.lcgadmin = fields[13]
    self.production = fields[14]
    self.pilot = fields[15]

    if ((re.match('^.*\d\d*$',self.name))      or
        (re.match('^.*\d\d*$',self.name3))     or
        (re.match('^.*\d\d*$',self.name5))     or
        (re.match('^.*\d\d*$',self.name8)) ):
      print('Warning: At least one of the name stubs of VO ' + self.blockname + ' has trailing digits' )

    if (self.lcgadmin != 'yes' and self.lcgadmin != 'no'):
      print('Warning: VO ' + self.blockname + ' has has bad value in lcgadmin field' )
    if (self.production != 'yes' and self.production != 'no'):
      print('Warning: VO ' + self.blockname + ' has has bad value in production field' )
    if (self.pilot != 'yes' and self.pilot != 'no'):
      print('Warning: VO ' + self.blockname + ' has has bad value in pilot field' )

    # ARC fields
    self.share = fields[16]

    # Sub groups
    extraGroups = (len(fields) - 17)/5
    for i in range (0,extraGroups):
      theName = fields[17+0]
      theUsers = fields[17+1]
      theFqdn = fields[17+2]
      theShare = fields[17+3]
      theGuid = fields[17+4]
      newSg = SubGroup ([theName,theUsers,theFqdn,theShare,theGuid])
      self.sg_list.append(newSg)

#--- class
class ConfReader:

  """ Reads the conf """

  # Constructor - set global values to null
  def __init__(self):
    self.start_uid = -1
    self.start_add_uid = -1

  # Read conf file
  def read_conf(self,vo_data_input,vo_list):

    # To find duplicate gids
    gidHash = {}

    # Conf file has several sections; 
    section = ''

    try:
      f = open(vo_data_input, 'r')
    except IOError:
      print('Cannot open input conf file'); sys.exit(1)

    # Read all the lines
    for line in f: 
      line = line.rstrip()
      line = re.sub('\#.*', '', line)

      # Omit blanks
      if re.match('^\s*$',line,0):
        continue

      # Track the section
      if re.match('^\s*\[general\]\s*$',line):
        section = 'general'
        continue
      if re.match('^\s*\[vos\]\s*$',line):
        section = 'vos'
        continue

      # Sorry - no unknown sections allowed.
      if (section != 'general')&(section != 'vos'):
        print('Odd conf file, check the section; is ' + section); sys.exit(1)

      if section == 'general':
        # Expect STARTUID STARTADDUID STARTSUBGROUPUID in this section
        matchObj = re.match('STARTUID\=(.*)',line)
        if matchObj:
          self.start_uid = int(matchObj.group(1).strip())

        matchObj = re.match('STARTADDUID\=(.*)',line)
        if matchObj:
          self.start_add_uid = int(matchObj.group(1).strip())

        matchObj = re.match('STARTSUBGROUPUID\=(.*)',line)
        if matchObj:
          self.start_sguid = int(matchObj.group(1).strip())

      elif section == 'vos':
        # Expect VO lines in this section

        # Break up the fields
        fields = re.split('\s*,\s*',line)

        # Sorry - this line has the wrong number of fields 
        base = len(fields) - 17
        if (base % 5 != 0):
          print('Odd conf file, check the fields; line is ' + line); sys.exit(1)

        # Make a new VO with these fields, and stick it in the list
        newVo = Vo(fields)

        # Validate
        if newVo.gid in gidHash.keys():
          print('Odd conf file, duplicate gid, check the fields; line is ' + line); sys.exit(1)
        gidHash[newVo.gid] = 1; 

        if newVo.pgid in gidHash.keys():
          print('Odd conf file, duplicate pgid, check the fields; line is ' + line); sys.exit(1)
        gidHash[newVo.pgid] = 1; 

        if newVo.sgid in gidHash.keys():
          print('Odd conf file, duplicate sgid, check the fields; line is ' + line); sys.exit(1)
        gidHash[newVo.sgid] = 1; 

        if (newVo.pilgid in gidHash.keys() and (newVo.pilgid != -1)):
          print('Odd conf file, duplicate pilgid, check the fields; line is ' + line); sys.exit(1)
        gidHash[newVo.pilgid] = 1; 

        vo_list.append(newVo)

    f.close()
    # vo_list.sort(key=lambda x: x.blockname);

    # Sorry - the STARTUID, STARTADDUID and at least one VO are needed
    if (len(vo_list) <= 0):
      print('Odd conf file, no VOs found'); sys.exit(1)

    if (self.start_uid < 0  ):
      print('Odd conf file, no STARTUID found'); sys.exit(1)

    if (self.start_add_uid < 0  ):
      print('Odd conf file, no STARTADDUID found'); sys.exit(1)

    if (self.start_sguid < 0  ):
      print('Odd conf file, no STARTSUBGROUPUID found'); sys.exit(1)

    #-------------------------------------
    # Now check if overlap
    uidHash = {}

    vouid = self.start_uid
    voadduid = self.start_add_uid

    # For over the VOs
    for vo in vo_list:

      # STD USERS
      nuid = vouid
      for jj in range(vo.users):
        if nuid in uidHash.keys(): 
          print('Odd conf file, duplicate uid, check lines around vo ' + vo.name + '\n'); sys.exit(1)
        uidHash[nuid] = 1
        nuid = nuid + 1
    
      # PRODUCTION USERS
      nuid = vouid + 300
      for jj in range(vo.prd ):
        if nuid in uidHash.keys():
          print('Odd conf file, duplicate uid, check lines around vo ' + vo.name + '\n'); sys.exit(1)
        uidHash[nuid] = 1
        nuid = nuid + 1

      # SGM USERS
      nuid = vouid + 400
      for jj in range(vo.sgm ):
        if nuid in uidHash.keys():
          print('Odd conf file, duplicate uid, check lines around vo ' + vo.name + '\n'); sys.exit(1)
        uidHash[nuid] = 1
        nuid = nuid + 1

      # PILOT USERS 
      nuid = voadduid
      for jj in range(vo.pil ):
        if nuid in uidHash.keys():
          print('Odd conf file, duplicate uid, check lines around vo ' + vo.name + ' or the STARTADDUID variable\n'); sys.exit(1)
        uidHash[nuid] = 1
        nuid = nuid + 1
 
      # set vouid and voadduid for next VO
      vouid = vouid + 500
      voadduid = voadduid + 250

#--- main 

# Get the options

options = {} 
initOptions(options)

# Set up a ConfReader, and read the vo_list, start_uid, start_add_uid , start_sguid

vo_list = []
conf_reader = ConfReader()
conf_reader.read_conf(options['vo_data_input'],vo_list)
start_uid = conf_reader.start_uid
start_add_uid = conf_reader.start_add_uid
start_sguid = conf_reader.start_sguid

# Set up a UsersConf, and write the users.conf

users_conf = UsersConf()
users_conf.write_users(vo_list,options['users_conf_output'],start_uid,start_add_uid, start_sguid)

# Set up a CondorConf, and write the condor.conf

condor_conf = CondorConf()
condor_conf.write_condor(vo_list,options['condor_conf_output'])

# Set up a GroupMapFile, and write the groupmapfile

groupmapfile = GroupMapFile()
groupmapfile.write_groupmapfile(vo_list,options['groupmap_output'])

# Set up a GridMapFile, and write the gridmapfile

gridmapfile = GridMapFile()
gridmapfile.write_gridmapfile(vo_list,options['gridmap_output'])

# Set up an Argus policy file for WNs

arguswn_conf = ArguswnConf()
arguswn_conf.write_arguswn(vo_list,options['arguswn_conf_output'])

#--- end