Migrate to Python 3
[collecty.git] / src / collecty / client.py
1 #!/usr/bin/python3
2 ###############################################################################
3 #                                                                             #
4 # collecty - A system statistics collection daemon for IPFire                 #
5 # Copyright (C) 2012 IPFire development team                                  #
6 #                                                                             #
7 # This program is free software: you can redistribute it and/or modify        #
8 # it under the terms of the GNU General Public License as published by        #
9 # the Free Software Foundation, either version 3 of the License, or           #
10 # (at your option) any later version.                                         #
11 #                                                                             #
12 # This program is distributed in the hope that it will be useful,             #
13 # but WITHOUT ANY WARRANTY; without even the implied warranty of              #
14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the               #
15 # GNU General Public License for more details.                                #
16 #                                                                             #
17 # You should have received a copy of the GNU General Public License           #
18 # along with this program.  If not, see <http://www.gnu.org/licenses/>.       #
19 #                                                                             #
20 ###############################################################################
21
22 import argparse
23 import dbus
24 import sys
25
26 from .constants import *
27 from .i18n import _
28
29 import logging
30 log = logging.getLogger("collectly.client")
31
32 class CollectyClient(object):
33         def __init__(self):
34                 self.bus = dbus.SystemBus()
35
36                 self.proxy = self.bus.get_object(BUS_DOMAIN, "/GraphGenerator")
37
38         def list_templates(self):
39                 templates = self.proxy.ListTemplates()
40
41                 return ["%s" % t for t in templates]
42
43         def list_templates_cli(self, ns):
44                 templates = self.list_templates()
45
46                 for t in sorted(templates):
47                         print(t)
48
49         def generate_graph(self, template_name, **kwargs):
50                 byte_array = self.proxy.GenerateGraph(template_name, kwargs,
51                         signature="sa{sv}")
52
53                 # Convert the byte array into a byte string again
54                 if byte_array:
55                         return "".join((chr(b) for b in byte_array))
56
57         def generate_graph_cli(self, ns):
58                 kwargs = {
59                         "format"    : ns.format,
60                         "object_id" : ns.object,
61                 }
62
63                 if ns.height or ns.width:
64                         kwargs.update({
65                                 "height" : ns.height or 0,
66                                 "width"  : ns.width or 0,
67                         })
68
69                 if ns.interval:
70                         kwargs["interval"] = ns.interval
71
72                 # Generate the graph image
73                 graph = self.generate_graph(ns.template, **kwargs)
74
75                 # Write file to disk
76                 with open(ns.filename, "wb") as f:
77                         f.write(graph)
78
79         def parse_cli(self, args):
80                 parser = argparse.ArgumentParser(prog="collecty-client")
81                 subparsers = parser.add_subparsers(help="sub-command help")
82
83                 # generate-graph
84                 parser_generate_graph = subparsers.add_parser("generate-graph",
85                         help=_("Generate a graph image"))
86                 parser_generate_graph.set_defaults(func=self.generate_graph_cli)
87                 parser_generate_graph.add_argument("--filename",
88                         help=_("filename"), required=True)
89                 parser_generate_graph.add_argument("--format",
90                         help=_("image format"), default=DEFAULT_IMAGE_FORMAT)
91                 parser_generate_graph.add_argument("--interval", help=_("interval"))
92                 parser_generate_graph.add_argument("--object",
93                         help=_("Object identifier"), default="default")
94                 parser_generate_graph.add_argument("--template",
95                         help=_("The graph template identifier"), required=True)
96
97                 # Dimensions
98                 parser_generate_graph.add_argument("--height", type=int, default=0,
99                         help=_("Height of the generated image"))
100                 parser_generate_graph.add_argument("--width", type=int, default=0,
101                         help=_("Width of the generated image"))
102
103                 # list-templates
104                 parser_list_templates = subparsers.add_parser("list-templates",
105                         help=_("Lists all graph templates"))
106                 parser_list_templates.set_defaults(func=self.list_templates_cli)
107
108                 return parser.parse_args(args)
109
110         def run_cli(self, args=None):
111                 args = self.parse_cli(args or sys.argv[1:])
112
113                 return args.func(args)