uploader.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430
  1. # -*- coding: utf-8 -*-
  2. # Copyright (C) 2015-2016 Peter Magnusson <peter@birchroad.net>
  3. """Main functionality for nodemcu-uploader"""
  4. import time
  5. import logging
  6. import hashlib
  7. import os
  8. import serial
  9. from .exceptions import CommunicationTimeout, DeviceNotFoundException
  10. from .exceptions import BadResponseException
  11. from .utils import default_port, system
  12. from .luacode import LUA_FUNCTIONS, DOWNLOAD_FILE, RECV_LUA, SEND_LUA, LIST_FILES, UART_SETUP, PRINT_FILE
  13. log = logging.getLogger(__name__) # pylint: disable=C0103
  14. __all__ = ['Uploader', 'default_port']
  15. SYSTEM = system()
  16. MINIMAL_TIMEOUT = 0.0001
  17. BLOCK_START = '\x01'
  18. NUL = '\x00'
  19. ACK = '\x06'
  20. class Uploader(object):
  21. """Uploader is the class for communicating with the nodemcu and
  22. that will allow various tasks like uploading files, formating the filesystem etc.
  23. """
  24. BAUD = 115200
  25. START_BAUD = 9600
  26. TIMEOUT = 5
  27. PORT = default_port()
  28. def __init__(self, port=PORT, baud=BAUD, start_baud=START_BAUD):
  29. log.info('opening port %s with %s baud', port, start_baud)
  30. if port == 'loop://':
  31. self._port = serial.serial_for_url(port, start_baud, timeout=Uploader.TIMEOUT)
  32. else:
  33. self._port = serial.Serial(port, start_baud, timeout=Uploader.TIMEOUT)
  34. self.start_baud = start_baud
  35. self.baud = baud
  36. # Keeps things working, if following conections are made:
  37. ## RTS = CH_PD (i.e reset)
  38. ## DTR = GPIO0
  39. self._port.setRTS(False)
  40. self._port.setDTR(False)
  41. def __sync():
  42. """Get in sync with LUA (this assumes that NodeMCU gets reset by the previous two lines)"""
  43. log.debug('getting in sync with LUA')
  44. self.__clear_buffers()
  45. try:
  46. self.__exchange(';') # Get a defined state
  47. self.__writeln('print("%sync%");')
  48. self.__expect('%sync%\r\n> ')
  49. except CommunicationTimeout:
  50. raise DeviceNotFoundException('Device not found or wrong port')
  51. __sync()
  52. if baud != start_baud:
  53. self.__set_baudrate(baud)
  54. # Get in sync again
  55. __sync()
  56. self.line_number = 0
  57. def __set_baudrate(self, baud):
  58. """setting baudrate if supported"""
  59. log.info('Changing communication to %s baud', baud)
  60. self.__writeln(UART_SETUP.format(baud=baud))
  61. # Wait for the string to be sent before switching baud
  62. time.sleep(0.1)
  63. try:
  64. self._port.setBaudrate(baud)
  65. except AttributeError:
  66. #pySerial 2.7
  67. self._port.baudrate = baud
  68. def __clear_buffers(self):
  69. """Clears the input and output buffers"""
  70. try:
  71. self._port.reset_input_buffer()
  72. self._port.reset_output_buffer()
  73. except AttributeError:
  74. #pySerial 2.7
  75. self._port.flushInput()
  76. self._port.flushOutput()
  77. def __expect(self, exp='> ', timeout=TIMEOUT):
  78. """will wait for exp to be returned from nodemcu or timeout"""
  79. timer = self._port.timeout
  80. #do NOT set timeout on Windows
  81. if SYSTEM != 'Windows':
  82. # Checking for new data every 100us is fast enough
  83. if self._port.timeout != MINIMAL_TIMEOUT:
  84. self._port.timeout = MINIMAL_TIMEOUT
  85. end = time.time() + timeout
  86. # Finish as soon as either exp matches or we run out of time (work like dump, but faster on success)
  87. data = ''
  88. while not data.endswith(exp) and time.time() <= end:
  89. data += self._port.read()
  90. log.debug('expect returned: `{0}`'.format(data))
  91. if time.time() > end:
  92. raise CommunicationTimeout('Timeout waiting for data', data)
  93. if not data.endswith(exp) and len(exp) > 0:
  94. raise BadResponseException('Bad response.', exp, data)
  95. if SYSTEM != 'Windows':
  96. self._port.timeout = timer
  97. return data
  98. def __write(self, output, binary=False):
  99. """write data on the nodemcu port. If 'binary' is True the debug log
  100. will show the intended output as hex, otherwise as string"""
  101. if not binary:
  102. log.debug('write: %s', output)
  103. else:
  104. log.debug('write binary: %s', ':'.join(x.encode('hex') for x in output))
  105. self._port.write(output)
  106. self._port.flush()
  107. def __writeln(self, output):
  108. """write, with linefeed"""
  109. self.__write(output + '\n')
  110. def __exchange(self, output, timeout=TIMEOUT):
  111. """Write output to the port and wait for response"""
  112. self.__writeln(output)
  113. self._port.flush()
  114. return self.__expect(timeout=timeout)
  115. def close(self):
  116. """restores the nodemcu to default baudrate and then closes the port"""
  117. try:
  118. if self.baud != self.start_baud:
  119. self.__set_baudrate(self.start_baud)
  120. self._port.flush()
  121. self.__clear_buffers()
  122. except serial.serialutil.SerialException:
  123. pass
  124. log.debug('closing port')
  125. self._port.close()
  126. def prepare(self):
  127. """
  128. This uploads the protocol functions nessecary to do binary
  129. chunked transfer
  130. """
  131. log.info('Preparing esp for transfer.')
  132. for func in LUA_FUNCTIONS:
  133. detected = self.__exchange('print({0})'.format(func))
  134. if detected.find('function:') == -1:
  135. break
  136. else:
  137. log.info('Preparation already done. Not adding functions again.')
  138. return True
  139. functions = RECV_LUA + '\n' + SEND_LUA
  140. data = functions.format(baud=self._port.baudrate)
  141. ##change any \r\n to just \n and split on that
  142. lines = data.replace('\r', '').split('\n')
  143. #remove some unneccesary spaces to conserve some bytes
  144. for line in lines:
  145. line = line.strip().replace(', ', ',').replace(' = ', '=')
  146. if len(line) == 0:
  147. continue
  148. resp = self.__exchange(line)
  149. #do some basic test of the result
  150. if ('unexpected' in resp) or ('stdin' in resp) or len(resp) > len(functions)+10:
  151. log.error('error when preparing "%s"', resp)
  152. return False
  153. return True
  154. def download_file(self, filename):
  155. """Download a file from device to local filesystem"""
  156. res = self.__exchange('send("{filename}")'.format(filename=filename))
  157. if ('unexpected' in res) or ('stdin' in res):
  158. log.error('Unexpected error downloading file: %s', res)
  159. raise Exception('Unexpected error downloading file')
  160. #tell device we are ready to receive
  161. self.__write('C')
  162. #we should get a NUL terminated filename to start with
  163. sent_filename = self.__expect(NUL).strip()
  164. log.info('receiveing ' + sent_filename)
  165. #ACK to start download
  166. self.__write(ACK, True)
  167. buf = ''
  168. data = ''
  169. chunk, buf = self.__read_chunk(buf)
  170. #read chunks until we get an empty which is the end
  171. while chunk != '':
  172. self.__write(ACK, True)
  173. data = data + chunk
  174. chunk, buf = self.__read_chunk(buf)
  175. return data
  176. def read_file(self, filename, destination=''):
  177. if not destination:
  178. destination = filename
  179. log.info('Transfering %s to %s', filename, destination)
  180. data = self.download_file(filename)
  181. with open(destination, 'w') as fil:
  182. fil.write(data)
  183. def write_file(self, path, destination='', verify='none'):
  184. """sends a file to the device using the transfer protocol"""
  185. filename = os.path.basename(path)
  186. if not destination:
  187. destination = filename
  188. log.info('Transfering %s as %s', path, destination)
  189. self.__writeln("recv()")
  190. res = self.__expect('C> ')
  191. if not res.endswith('C> '):
  192. log.error('Error waiting for esp "%s"', res)
  193. return
  194. log.debug('sending destination filename "%s"', destination)
  195. self.__write(destination + '\x00', True)
  196. if not self.__got_ack():
  197. log.error('did not ack destination filename')
  198. return
  199. fil = open(path, 'rb')
  200. content = fil.read()
  201. fil.close()
  202. log.debug('sending %d bytes in %s', len(content), filename)
  203. pos = 0
  204. chunk_size = 128
  205. while pos < len(content):
  206. rest = len(content) - pos
  207. if rest > chunk_size:
  208. rest = chunk_size
  209. data = content[pos:pos+rest]
  210. if not self.__write_chunk(data):
  211. resp = self.__expect()
  212. log.error('Bad chunk response "%s" %s', resp, ':'.join(x.encode('hex') for x in resp))
  213. return
  214. pos += chunk_size
  215. log.debug('sending zero block')
  216. #zero size block
  217. self.__write_chunk('')
  218. if verify == 'raw':
  219. log.info('Verifying...')
  220. data = self.download_file(destination)
  221. if content != data:
  222. log.error('Verification failed.')
  223. else:
  224. log.info('Verification successfull. Contents are identical.')
  225. elif verify == 'sha1':
  226. #Calculate SHA1 on remote file. Extract just hash from result
  227. data = self.__exchange('shafile("'+destination+'")').splitlines()[1]
  228. log.info('Remote SHA1: %s', data)
  229. #Calculate hash of local data
  230. filehashhex = hashlib.sha1(content.encode()).hexdigest()
  231. log.info('Local SHA1: %s', filehashhex)
  232. if data != filehashhex:
  233. log.error('Verification failed.')
  234. else:
  235. log.info('Verification successfull. Checksums match')
  236. elif verify != 'none':
  237. raise Exception(verify + ' is not a valid verification method.')
  238. def exec_file(self, path):
  239. """execute the lines in the local file 'path'"""
  240. filename = os.path.basename(path)
  241. log.info('Execute %s', filename)
  242. fil = open(path, 'r')
  243. res = '> '
  244. for line in fil:
  245. line = line.rstrip('\r\n')
  246. retlines = (res + self.__exchange(line)).splitlines()
  247. # Log all but the last line
  248. res = retlines.pop()
  249. for lin in retlines:
  250. log.info(lin)
  251. # last line
  252. log.info(res)
  253. fil.close()
  254. def __got_ack(self):
  255. """Returns true if ACK is received"""
  256. log.debug('waiting for ack')
  257. res = self._port.read(1)
  258. log.debug('ack read %s', res.encode('hex'))
  259. return res == '\x06' #ACK
  260. def write_lines(self, data):
  261. lines = data.replace('\r', '').split('\n')
  262. for line in lines:
  263. self.__exchange(line)
  264. return
  265. def __write_chunk(self, chunk):
  266. """formats and sends a chunk of data to the device according
  267. to transfer protocol"""
  268. log.debug('writing %d bytes chunk', len(chunk))
  269. data = BLOCK_START + chr(len(chunk)) + chunk
  270. if len(chunk) < 128:
  271. padding = 128 - len(chunk)
  272. log.debug('pad with %d characters', padding)
  273. data = data + (' ' * padding)
  274. log.debug("packet size %d", len(data))
  275. self.__write(data)
  276. self._port.flush()
  277. return self.__got_ack()
  278. def __read_chunk(self, buf):
  279. """Read a chunk of data"""
  280. log.debug('reading chunk')
  281. timeout = self._port.timeout
  282. if SYSTEM != 'Windows':
  283. # Checking for new data every 100us is fast enough
  284. if self._port.timeout != MINIMAL_TIMEOUT:
  285. self._port.timeout = MINIMAL_TIMEOUT
  286. end = time.time() + timeout
  287. while len(buf) < 130 and time.time() <= end:
  288. buf = buf + self._port.read()
  289. if buf[0] != BLOCK_START or len(buf) < 130:
  290. print 'buffer size:', len(buf)
  291. log.debug('buffer binary: %s ', ':'.join(x.encode('hex') for x in buf))
  292. raise Exception('Bad blocksize or start byte')
  293. if SYSTEM != 'Windows':
  294. self._port.timeout = timeout
  295. chunk_size = ord(buf[1])
  296. data = buf[2:chunk_size+2]
  297. buf = buf[130:]
  298. return (data, buf)
  299. def file_list(self):
  300. """list files on the device"""
  301. log.info('Listing files')
  302. res = self.__exchange(LIST_FILES)
  303. log.info(res)
  304. return res
  305. def file_do(self, filename):
  306. """Execute a file on the device using 'do'"""
  307. log.info('Executing '+filename)
  308. res = self.__exchange('dofile("'+filename+'")')
  309. log.info(res)
  310. return res
  311. def file_format(self):
  312. """Formats device filesystem"""
  313. log.info('Formating, can take up to 1 minute...')
  314. res = self.__exchange('file.format()', timeout=60)
  315. if 'format done' not in res:
  316. log.error(res)
  317. else:
  318. log.info(res)
  319. return res
  320. def file_print(self, filename):
  321. """Prints a file on the device to console"""
  322. log.info('Printing ' + filename)
  323. res = self.__exchange(PRINT_FILE.format(filename=filename))
  324. log.info(res)
  325. return res
  326. def node_heap(self):
  327. """Show device heap size"""
  328. log.info('Heap')
  329. res = self.__exchange('print(node.heap())')
  330. log.info(res)
  331. return res
  332. def node_restart(self):
  333. """Restarts device"""
  334. log.info('Restart')
  335. res = self.__exchange('node.restart()')
  336. log.info(res)
  337. return res
  338. def file_compile(self, path):
  339. """Compiles a file specified by path on the device"""
  340. log.info('Compile '+path)
  341. cmd = 'node.compile("%s")' % path
  342. res = self.__exchange(cmd)
  343. log.info(res)
  344. return res
  345. def file_remove(self, path):
  346. """Removes a file on the device"""
  347. log.info('Remove '+path)
  348. cmd = 'file.remove("%s")' % path
  349. res = self.__exchange(cmd)
  350. log.info(res)
  351. return res