#!/usr/bin/env python
# -*- coding:utf-8 -*- 
 
 
import subprocess
import json
import sys
import argparse
 
def R(s):
 return "%s[31;2m%s%s[0m"%(chr(27), s, chr(27))
 
def get_present_sgid(vmid):
 descmd = '/usr/bin/qcloudcli dfw DescribeSecurityGroups --instanceId ' + vmid.strip()
 p = subprocess.Popen(descmd, shell=True, stdout=subprocess.PIPE) 
 output = p.communicate()[0]
 res = json.loads(output)
 sgid = []
 for d in res['data']:
  sid = d['sgId']
  sgid.append(str(sid))
 return sgid
 
def make_json(vmid,sgid):
 pdata = {}
 pdata["instanceId"] = vmid
 pdata["sgIds"] = sgid
 pjson = json.dumps(pdata)
 return pjson
 
def add_sgid(vmfile,newsid):
 fi = open(vmfile)
 for v in fi:
  v = v.strip()
  res = get_present_sgid(v)
  print res
  res.append(newsid)
  pjson = make_json(v,res)
  modcmd = 'qcloudcli dfw ModifySecurityGroupsOfInstance --instanceSet ' + "'[" + pjson+ "]'"
  p = subprocess.Popen(modcmd, shell=True, stdout=subprocess.PIPE)
  output = p.communicate()[0]
  print output
 
def remove_sgid(vmfile,newsid):
 fi = open(vmfile)
 for v in fi:
  v = v.strip()
  res = get_present_sgid(v)
  res.remove(newsid)
  pjson = make_json(v,res)
  modcmd = 'qcloudcli dfw ModifySecurityGroupsOfInstance --instanceSet ' + "'[" + pjson+ "]'"
  p = subprocess.Popen(modcmd, shell=True, stdout=subprocess.PIPE)
  output = p.communicate()[0]
  #print output
 
 
if __name__ == "__main__":
 parser=argparse.ArgumentParser(description='change sgid', usage='%(prog)s [options]')
 parser.add_argument('-f','--file', nargs='?', dest='filehost', help='vmidfile')
 parser.add_argument('-g','--sgid', nargs='?', dest='sgid', help='sgid')
 parser.add_argument('-m','--method', nargs='?', dest='method', help='Methods only support to add or remove')
 if len(sys.argv)==1:
  parser.print_help()
 else:
  args=parser.parse_args()
  if args.filehost is not None and args.sgid is not None and args.method is not None:
   if args.method == 'add':
    add_sgid(args.filehost, args.sgid)
   elif args.method == 'remove':
    remove_sgid(args.filehost, args.sgid)
   else:
    print R('Methods only support to add or remove')
  else:
   print R('Error format, please see the usage:')
   parser.print_help()