summaryrefslogtreecommitdiffstats
path: root/scripts/automation/trex_control_plane/client_utils/jsonrpc_client.py
blob: aff6b36ece2e0a7d0888e0b24140d484b433fe58 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
#!/router/bin/python

import outer_packages
import zmq
import json
import general_utils
import re
from time import sleep

class bcolors:
    BLUE = '\033[94m'
    GREEN = '\033[32m'
    YELLOW = '\033[93m'
    RED = '\033[31m'
    MAGENTA = '\033[35m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


class JsonRpcClient(object):

    def __init__ (self, default_server, default_port):
        self.verbose = False
        self.connected = False

        # default values
        self.port   = default_port
        self.server = default_server
        self.id_gen = general_utils.random_id_gen()

    def get_connection_details (self):
        rc = {}
        rc['server'] = self.server
        rc['port']   = self.port

        return rc

    def pretty_json (self, json_str, use_colors = True):
        pretty_str = json.dumps(json.loads(json_str), indent = 4, separators=(',', ': '), sort_keys = True)

        if not use_colors:
            return pretty_str

        try:
            # int numbers
            pretty_str = re.sub(r'([ ]*:[ ]+)(\-?[1-9][0-9]*[^.])',r'\1{0}\2{1}'.format(bcolors.BLUE, bcolors.ENDC), pretty_str)
            # float
            pretty_str = re.sub(r'([ ]*:[ ]+)(\-?[1-9][0-9]*\.[0-9]+)',r'\1{0}\2{1}'.format(bcolors.MAGENTA, bcolors.ENDC), pretty_str)
            # strings
            
            pretty_str = re.sub(r'([ ]*:[ ]+)("[^"]*")',r'\1{0}\2{1}'.format(bcolors.RED, bcolors.ENDC), pretty_str)
            pretty_str = re.sub(r"('[^']*')", r'{0}\1{1}'.format(bcolors.MAGENTA, bcolors.RED), pretty_str)
        except :
            pass

        return pretty_str

    def verbose_msg (self, msg):
        if not self.verbose:
            return

        print "[verbose] " + msg


    def create_jsonrpc_v2 (self, method_name, params = {}, id = None):
        msg = {}
        msg["jsonrpc"] = "2.0"
        msg["method"]  = method_name

        msg["params"] = params

        msg["id"] = id

        return json.dumps(msg)

    def invoke_rpc_method (self, method_name, params = {}, block = False):
        rc, msg = self._invoke_rpc_method(method_name, params, block)
        if not rc:
            self.disconnect()

        return rc, msg

    def _invoke_rpc_method (self, method_name, params = {}, block = False):
        if not self.connected:
            return False, "Not connected to server"

        id = self.id_gen.next()
        msg = self.create_jsonrpc_v2(method_name, params, id = id)

        self.verbose_msg("Sending Request To Server:\n\n" + self.pretty_json(msg) + "\n")

        if block:
            self.socket.send(msg)
        else:
            try:
                self.socket.send(msg, flags = zmq.NOBLOCK)
            except zmq.error.ZMQError as e:
                return False, "Failed To Get Send Message"

        got_response = False

        if block:
            response = self.socket.recv()
            got_response = True
        else:
            for i in xrange(0 ,10):
                try:
                    response = self.socket.recv(flags = zmq.NOBLOCK)
                    got_response = True
                    break
                except zmq.Again:
                    sleep(0.2)

        if not got_response:
            return False, "Failed To Get Server Response"

        self.verbose_msg("Server Response:\n\n" + self.pretty_json(response) + "\n")

        # decode
        response_json = json.loads(response)

        if (response_json.get("jsonrpc") != "2.0"):
            return False, "Malfromed Response ({0})".format(str(response))

        if (response_json.get("id") != id):
            return False, "Server Replied With Bad ID ({0})".format(str(response))

        # error reported by server
        if ("error" in response_json):
            return True, response_json["error"]["message"]

        # if no error there should be a result
        if ("result" not in response_json):
            return False, "Malfromed Response ({0})".format(str(response))

        return True, response_json["result"]


    def ping_rpc_server(self):

        return self.invoke_rpc_method("ping", block = False)

    def get_rpc_server_status (self):
        return self.invoke_rpc_method("get_status")

    def query_rpc_server(self):
        return self.invoke_rpc_method("get_supported_cmds")


    def set_verbose(self, mode):
        self.verbose = mode

    def disconnect (self):
        if self.connected:
            self.socket.close(linger = 0)
            self.context.destroy(linger = 0)
            self.connected = False
            return True, ""
        else:
            return False, "Not connected to server"

    def connect(self, server = None, port = None):
        if self.connected:
            self.disconnect()

        self.context = zmq.Context()

        self.server = (server if server else self.server)
        self.port = (port if port else self.port)

        #  Socket to talk to server
        self.transport = "tcp://{0}:{1}".format(self.server, self.port)

        print "\nConnecting To RPC Server On {0}".format(self.transport)

        self.socket = self.context.socket(zmq.REQ)
        try:
            self.socket.connect(self.transport)
        except zmq.error.ZMQError as e:
            return False, "ZMQ Error: Bad server or port name: " + str(e)


        self.connected = True

        # ping the server
        rc, err = self.ping_rpc_server()
        if not rc:
            self.disconnect()
            return rc, err

        return True, ""


    def reconnect(self):
        # connect using current values
        return self.connect()

        if not self.connected:
            return False, "Not connected to server"

        # reconnect
        return self.connect(self.server, self.port)


    def is_connected(self):
        return self.connected


    def __del__(self):
        print "Shutting down RPC client\n"
        self.context.destroy(linger=0)

if __name__ == "__main__":
    pass