Changed 'import auth' statements to use plc.py or monitorconfig.py
[monitor.git] / nodegroups.py
1 #!/usr/bin/python
2
3 # This script is used to manipulate the operational state of nodes in
4 # different node groups.  These are basically set operations on nodes via the
5 # PLC api.
6
7 # Take the ng name as an argument....
8 # optionally, 
9 #  * restart them all.
10 #  * Set some or all in the set to rins.
11 #  * get a list of nodes in the Alpha nodegroup.
12
13 # Given a nodelist, it could tag each one with a nodegroup name.
14 #  * 
15
16 import plc
17 api = plc.getAuthAPI()
18
19 from optparse import OptionParser
20 from sets import Set
21 from nodequery import verify,query_to_dict,node_select
22
23 from nodecommon import *
24 import database
25
26 def main():
27         from config import config
28         fb = database.dbLoad("findbad")
29
30         parser = OptionParser()
31         parser.set_defaults(nodegroup="Alpha",
32                                                 node=None,
33                                                 nodelist=None,
34                                                 list=True,
35                                                 add=False,
36                         nocolor=False,
37                                                 notng=False,
38                                                 delete=False,
39                                                 nodeselect=None,
40                                                 )
41         parser.add_option("", "--not", dest="notng", action="store_true", 
42                                                 help="All nodes NOT in nodegroup.")
43         parser.add_option("", "--nodegroup", dest="nodegroup", metavar="NodegroupName",
44                                                 help="Specify a nodegroup to perform actions on")
45         parser.add_option("", "--nodeselect", dest="nodeselect", metavar="querystring",
46                                                 help="Specify a query to perform on findbad db")
47         parser.add_option("", "--site", dest="site", metavar="site name",
48                                                 help="Specify a site to view node status")
49
50         parser.add_option("", "--nocolor", dest="nocolor", action="store_true", 
51                                                 help="Enable color")
52         parser.add_option("", "--list", dest="list", action="store_true", 
53                                                 help="List all nodes in the given nodegroup")
54         parser.add_option("", "--add", dest="add", action="store_true", 
55                                                 help="Add nodes to the given nodegroup")
56         parser.add_option("", "--delete", dest="delete", action="store_true", 
57                                                 help="Delete nodes from the given nodegroup")
58         parser.add_option("", "--node", dest="node", metavar="nodename.edu", 
59                                                 help="A single node name to add to the nodegroup")
60         parser.add_option("", "--nodelist", dest="nodelist", metavar="list.txt", 
61                                                 help="Use all nodes in the given file for operation.")
62         config = config(parser)
63         config.parse_args()
64
65         # COLLECT nodegroups, nodes and node lists
66         if config.node or config.nodelist:
67                 if config.node: 
68                         hostlist = [ config.node ] 
69                 else: 
70                         hostlist = config.getListFromFile(config.nodelist)
71
72                 # NOTE: preserve order given in file.  Otherwise, return values are not in order
73                 # given to GetNodes
74                 nodelist = []
75                 for h in hostlist:
76                         nodelist += api.GetNodes(h)
77
78                 #nodelist = api.GetNodes(hostlist)
79                 group_str = "Given"
80
81         elif config.site:
82                 site = api.GetSites(config.site)
83                 if len (site) > 0:
84                         site = site[0]
85                         nodelist = api.GetNodes(site['node_ids'])
86                 else:
87                         nodelist = []
88
89                 group_str = config.site
90
91         elif config.nodeselect:
92                 hostlist = node_select(config.nodeselect)
93                 nodelist = api.GetNodes(hostlist)
94
95                 group_str = "selection"
96                 
97         else:
98                 ng = api.GetNodeGroups({'name' : config.nodegroup})
99                 nodelist = api.GetNodes(ng[0]['node_ids'])
100
101                 group_str = config.nodegroup
102
103         if config.notng:
104                 # Get nodegroup nodes
105                 ng_nodes = nodelist
106
107                 # Get all nodes
108                 all_nodes = api.GetNodes({'peer_id': None})
109                 
110                 # remove ngnodes from all node list
111                 ng_list = [ x['hostname'] for x in ng_nodes ]
112                 all_list = [ x['hostname'] for x in all_nodes ]
113                 not_ng_nodes = Set(all_list) - Set(ng_list)
114
115                 # keep each node that *is* in the not_ng_nodes set
116                 nodelist = filter(lambda x : x['hostname'] in not_ng_nodes, all_nodes)
117
118         hostnames = [ n['hostname'] for n in nodelist ]
119
120         # commands:
121
122         if config.add and config.nodegroup:
123                 for node in hostnames:
124                         print "Adding %s to %s nodegroup" % (node, config.nodegroup)
125                         api.AddNodeToNodeGroup(node, config.nodegroup)
126
127         elif config.delete:
128                 for node in hostnames:
129                         print "Deleting %s from %s nodegroup" % (node, config.nodegroup)
130                         api.DeleteNodeFromNodeGroup(node, config.nodegroup)
131
132         elif config.list:
133                 print " ---- Nodes in the %s Node Group ----" % group_str
134                 i = 1
135                 for node in nodelist:
136                         print "%-2d" % i, 
137                         print nodegroup_display(node, fb, config)
138                         i += 1
139
140         else:
141                 print "no other options supported."
142
143 if __name__ == "__main__":
144         try:
145                 main()
146         except IOError:
147                 pass