1#!/usr/bin/env python3
2#
3#   Copyright 2019 - The Android Open Source Project
4#
5#   Licensed under the Apache License, Version 2.0 (the "License");
6#   you may not use this file except in compliance with the License.
7#   You may obtain a copy of the License at
8#
9#           http://www.apache.org/licenses/LICENSE-2.0
10#
11#   Unless required by applicable law or agreed to in writing, software
12#   distributed under the License is distributed on an "AS IS" BASIS,
13#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14#   See the License for the specific language governing permissions and
15#   limitations under the License.
16"""Python module for Abstract Instrument Library."""
17
18import socket
19import requests
20from acts import logger
21
22
23class SocketInstrumentError(Exception):
24    """Abstract Instrument Error Class, via Socket and SCPI."""
25
26    def __init__(self, error, command=None):
27        """Init method for Socket Instrument Error.
28
29        Args:
30            error: Exception error.
31            command: Additional information on command,
32                Type, Str.
33        """
34        super(SocketInstrumentError, self).__init__(error)
35        self._error_code = error
36        self._error_message = self._error_code
37        if command is not None:
38            self._error_message = 'Command {} returned the error: {}.'.format(
39                repr(command), repr(self._error_message))
40
41    def __str__(self):
42        return self._error_message
43
44
45class SocketInstrument(object):
46    """Abstract Instrument Class, via Socket and SCPI."""
47
48    def __init__(self, ip_addr, ip_port):
49        """Init method for Socket Instrument.
50
51        Args:
52            ip_addr: IP Address.
53                Type, str.
54            ip_port: TCPIP Port.
55                Type, str.
56        """
57        self._socket_timeout = 120
58        self._socket_buffer_size = 1024
59
60        self._ip_addr = ip_addr
61        self._ip_port = ip_port
62
63        self._escseq = '\n'
64        self._codefmt = 'utf-8'
65
66        self._logger = logger.create_tagged_trace_logger(
67            '%s:%s' % (self._ip_addr, self._ip_port))
68
69        self._socket = None
70
71    def _connect_socket(self):
72        """Init and Connect to socket."""
73        try:
74            self._socket = socket.create_connection(
75                (self._ip_addr, self._ip_port), timeout=self._socket_timeout)
76
77            infmsg = 'Opened Socket connection to {}:{} with handle {}.'.format(
78                repr(self._ip_addr), repr(self._ip_port), repr(self._socket))
79            self._logger.debug(infmsg)
80
81        except socket.timeout:
82            errmsg = 'Socket timeout while connecting to instrument.'
83            self._logger.exception(errmsg)
84            raise SocketInstrumentError(errmsg)
85
86        except socket.error:
87            errmsg = 'Socket error while connecting to instrument.'
88            self._logger.exception(errmsg)
89            raise SocketInstrumentError(errmsg)
90
91    def _send(self, cmd):
92        """Send command via Socket.
93
94        Args:
95            cmd: Command to send,
96                Type, Str.
97        """
98        if not self._socket:
99            self._logger.warning('Socket instrument is not connected')
100            self._connect_socket()
101
102        cmd_es = cmd + self._escseq
103
104        try:
105            self._socket.sendall(cmd_es.encode(self._codefmt))
106            self._logger.debug('Sent %r to %r:%r.', cmd, self._ip_addr,
107                               self._ip_port)
108
109        except socket.timeout:
110            errmsg = ('Socket timeout while sending command {} '
111                      'to instrument.').format(repr(cmd))
112            self._logger.exception(errmsg)
113            raise SocketInstrumentError(errmsg)
114
115        except socket.error:
116            errmsg = ('Socket error while sending command {} '
117                      'to instrument.').format(repr(cmd))
118            self._logger.exception(errmsg)
119            raise SocketInstrumentError(errmsg)
120
121        except Exception as err:
122            errmsg = ('Error {} while sending command {} '
123                      'to instrument.').format(repr(cmd), repr(err))
124            self._logger.exception(errmsg)
125            raise
126
127    def _recv(self):
128        """Receive response via Socket.
129
130        Returns:
131            resp: Response from Instrument via Socket,
132                Type, Str.
133        """
134        if not self._socket:
135            self._logger.warning('Socket instrument is not connected')
136            self._connect_socket()
137
138        resp = ''
139
140        try:
141            while True:
142                resp_tmp = self._socket.recv(self._socket_buffer_size)
143                resp_tmp = resp_tmp.decode(self._codefmt)
144                resp += resp_tmp
145                if len(resp_tmp) < self._socket_buffer_size:
146                    break
147
148        except socket.timeout:
149            errmsg = 'Socket timeout while receiving response from instrument.'
150            self._logger.exception(errmsg)
151            raise SocketInstrumentError(errmsg)
152
153        except socket.error:
154            errmsg = 'Socket error while receiving response from instrument.'
155            self._logger.exception(errmsg)
156            raise SocketInstrumentError(errmsg)
157
158        except Exception as err:
159            errmsg = ('Error {} while receiving response '
160                      'from instrument').format(repr(err))
161            self._logger.exception(errmsg)
162            raise
163
164        resp = resp.rstrip(self._escseq)
165
166        self._logger.debug('Received %r from %r:%r.', resp, self._ip_addr,
167                           self._ip_port)
168
169        return resp
170
171    def _close_socket(self):
172        """Close Socket Instrument."""
173        if not self._socket:
174            return
175
176        try:
177            self._socket.shutdown(socket.SHUT_RDWR)
178            self._socket.close()
179            self._socket = None
180            self._logger.debug('Closed Socket Instrument %r:%r.',
181                               self._ip_addr, self._ip_port)
182
183        except Exception as err:
184            errmsg = 'Error {} while closing instrument.'.format(repr(err))
185            self._logger.exception(errmsg)
186            raise
187
188    def _query(self, cmd):
189        """query instrument via Socket.
190
191        Args:
192            cmd: Command to send,
193                Type, Str.
194
195        Returns:
196            resp: Response from Instrument via Socket,
197                Type, Str.
198        """
199        self._send(cmd + ';*OPC?')
200        resp = self._recv()
201        return resp
202
203
204class RequestInstrument(object):
205    """Abstract Instrument Class, via Request."""
206
207    def __init__(self, ip_addr):
208        """Init method for request instrument.
209
210        Args:
211            ip_addr: IP Address.
212                Type, Str.
213        """
214        self._request_timeout = 120
215        self._request_protocol = 'http'
216        self._ip_addr = ip_addr
217        self._escseq = '\r\n'
218
219        self._logger = logger.create_tagged_trace_logger(self._ip_addr)
220
221    def _query(self, cmd):
222        """query instrument via request.
223
224        Args:
225            cmd: Command to send,
226                Type, Str.
227
228        Returns:
229            resp: Response from Instrument via request,
230                Type, Str.
231        """
232        request_cmd = '{}://{}/{}'.format(self._request_protocol,
233                                          self._ip_addr, cmd)
234        resp_raw = requests.get(request_cmd, timeout=self._request_timeout)
235
236        resp = resp_raw.text
237        for char_del in self._escseq:
238            resp = resp.replace(char_del, '')
239
240        self._logger.debug('Sent %r to %r, and get %r.', cmd, self._ip_addr,
241                           resp)
242
243        return resp
244