8000 WebSocket Client class for Python · micropython/micropython@3a9aec1 · GitHub
[go: up one dir, main page]
Skip to content

Commit 3a9aec1

Browse files
CarglglzCarglglz
authored
WebSocket Client class for Python
Simple class to send commands and receive output
1 parent 8998d0f commit 3a9aec1

File tree

1 file changed

+170
-0
lines changed

1 file changed

+170
-0
lines changed

‎tools/websocket_client.py

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
1+
import ast
2+
import ws_client_handshake as wsclient
3+
4+
class BASE_WS_DEVICE:
5+
def __init__(self, target, password, init=False, ssl=False, auth=False, capath='.'):
6+
self.ws = None
7+
self.ip = target
8+
self.pswd = password
9+
self.port = 8266
10+
self.bytes_sent = 0
11+
self.buff = b''
12+
self.raw_buff = b''
13+
self.prompt = b'>>> '
14+
self.response = ''
15+
self._kbi = '\x03'
16+
self._banner = '\x02'
17+
self._reset = '\x04'
18+
self._traceback = b'Traceback (most recent call last):'
19+
self._flush = b''
20+
self.output = None
21+
self.platform = None
22+
self.connected = False
23+
if init:
24+
if not ssl:
25+
self.ws = wsclient.connect('ws://{}:{}'.format(self.ip, self.port), self.pswd)
26+
else:
27+
self.port = 8833
28+
self.ws = wsclient.connect('wss://{}:{}'.format(self.ip, self.port), self.pswd, auth=auth, capath=capath)
29+
self.connected = True
30+
31+
def open_wconn(self, ssl=False, auth=False, capath='.'):
32+
if not ssl:
33+
self.ws = wsclient.connect('ws://{}:{}'.format(self.ip, self.port), self.pswd)
34+
else:
35+
self.port = 8833
36+
self.ws = wsclient.connect('wss://{}:{}'.format(self.ip, self.port), self.pswd, auth=auth, capath=capath)
37+
self.connected = True
38+
39+
def close_wconn(self):
40+
self.ws.close()
41+
self.connected = False
42+
43+
def write(self, cmd):
44+
n_bytes = len(bytes(cmd, 'utf-8'))
45+
self.ws.send(cmd)
46+
return n_bytes
47+
48+
def read_all(self):
49+
self.ws.sock.settimeout(None)
50+
try:
51+
self.raw_buff = b''
52+
while self.prompt not in self.raw_buff:
53+
fin, opcode, data = self.ws.read_frame()
54+
self.raw_buff += data
55+
56+
return self.raw_buff
57+
except socket.timeout as e:
58+
return self.raw_buff
59+
60+
def flush(self):
61+
self.ws.sock.settimeout(0.01)
62+
self._flush = b''
63+
while True:
64+
try:
65+
fin, opcode, data = self.ws.read_frame()
66+
self._flush += data
67+
except socket.timeout as e:
68+
break
69+
except protocol.NoDataException as e:
70+
break
71+
72+
def wr_cmd(self, cmd, silent=False, rtn=True, rtn_resp=False, long_string=False):
73+
self.output = None
74+
self.response = ''
75+
self.buff = b''
76+
self.flush()
77+
self.bytes_sent = self.write(cmd+'\r')
78+
# time.sleep(0.1)
79+
# self.buff = self.read_all()[self.bytes_sent:]
80+
self.buff = self.read_all()
81+
if self.buff == b'':
82+
# time.sleep(0.1)
83+
self.buff = self.read_all()
84+
# print(self.buff)
85+
# filter command
86+
cmd_filt = bytes(cmd + '\r\n', 'utf-8')
87+
self.buff = self.buff.replace(cmd_filt, b'', 1)
88+
if self._traceback in self.buff:
89+
long_string = True
90+
if long_string:
91+
self.response = self.buff.replace(b'\r', b'').replace(b'\r\n>>> ', b'').replace(b'>>> ', b'').decode()
92+
else:
93+
self.response = self.buff.replace(b'\r\n', b'').replace(b'\r\n>>> ', b'').replace(b'>>> ', b'').decode()
94+
if not silent:
95+
if self.response != '\n' and self.response != '':
96+
print(self.response)
97+
else:
98+
self.response = ''
99+
if rtn:
100+
self.get_output()
101+
if self.output == '\n' and self.output == '':
102+
self.output = None
103+
if self.output is None:
104+
if self.response != '' and self.response != '\n':
105+
self.output = self.response
106+
if rtn_resp:
107+
return self.output
108+
109+
def cmd(self, cmd, silent=False, rtn=False, ssl=False):
110+
self.open_wconn(ssl=ssl, auth=True)
111+
self.wr_cmd(cmd, silent=True)
112+
self.close_wconn()
113+
self.get_output()
114+
if not silent:
115+
print(self.response)
116+
if rtn:
117+
return self.output
118+
119+
def reset(self, silent=False, ssl=False):
120+
if not silent:
121+
print('Rebooting device...')
122+
if self.connected:
123+
self.bytes_sent = self.write(self._reset)
124+
self.close_wconn()
125+
time.sleep(1)
126+
while True:
127+
try:
128+
self.open_wconn()
129+
self.wr_cmd(self._banner, silent=True)
130+
break
131+
except Exception as e:
132+
time.sleep(0.5)
133+
if not silent:
134+
print('Done!')
135+
else:
136+
self.open_wconn(ssl=ssl, auth=True)
137+
self.bytes_sent = self.write(self._reset)
138+
self.close_wconn()
139+
if not silent:
140+
print('Done!')
141+
142+
def kbi(self, silent=True):
143+
if self.connected:
144+
self.wr_cmd(self._kbi, silent=silent)
145+
else:
146+
self.cmd(self._kbi, silent=silent)
147+
148+
def banner(self):
149+
self.wr_cmd(self._banner, silent=True, long_string=True)
150+
print(self.response.replace('\n\n', '\n'))
151+
152+
def get_output(self):
153+
try:
154+
self.output = ast.literal_eval(self.response)
155+
except Exception as e:
156+
if 'bytearray' in self.response:
157+
try:
158+
self.output = bytearray(ast.literal_eval(
159+
self.response.strip().split('bytearray')[1]))
160+
except Exception as e:
161+
pass
162+
else:
163+
if 'array' in self.response:
164+
try:
165+
arr = ast.literal_eval(
166+
self.response.strip().split('array')[1])
167+
self.output = array(arr[0], arr[1])
168+
except Exception as e:
169+
pass
170+
pass

0 commit comments

Comments
 (0)
0