primadiag-pybash/picowatch_d/picowatch.py

#!/usr/bin/env python

# Primadiag SAS - Paris
# Author: Gino D.
# Copyright (c) 2023 Primadiag
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE

import os
import sys
import time
import json
import serial
import binascii
import textwrap

from watchdog.observers import Observer
from watchdog.events import PatternMatchingEventHandler
from typing import Dict, List, Optional, Tuple


BUFFER_SIZE: int = 256


# class TelnetToSerial:
#     def __init__(self, ip, user, password, read_timeout=None):
#         import telnetlib
#         self.tn = telnetlib.Telnet(ip, timeout=15)
#         self.read_timeout = read_timeout
#         if b'Login as:' in self.tn.read_until(b'Login as:', timeout=read_timeout):
#             self.tn.write(bytes(user, 'ascii') + b"\r\n")

#             if b'Password:' in self.tn.read_until(b'Password:', timeout=read_timeout):
#                 # needed because of internal implementation details of the telnet server
#                 time.sleep(0.2)
#                 self.tn.write(bytes(password, 'ascii') + b"\r\n")

#                 if b'for more information.' in self.tn.read_until(b'Type "help()" for more information.', timeout=read_timeout):
#                     # login succesful
#                     from collections import deque
#                     self.fifo = deque()
#                     return

#         raise Exception('Failed to establish a telnet connection with the board')

#     def __del__(self):
#         self.close()

#     def close(self):
#         try:
#             self.tn.close()
#         except:
#             # the telnet object might not exist yet, so ignore this one
#             pass

#     def read(self, size=1):
#         while len(self.fifo) < size:
#             timeout_count = 0
#             data = self.tn.read_eager()
#             if len(data):
#                 self.fifo.extend(data)
#                 timeout_count = 0
#             else:
#                 time.sleep(0.25)
#                 if self.read_timeout is not None and timeout_count > 4 * self.read_timeout:
#                     break
#                 timeout_count += 1

#         data = b''
#         while len(data) < size and len(self.fifo) > 0:
#             data += bytes([self.fifo.popleft()])
#         return data

#     def write(self, data):
#         self.tn.write(data)
#         return len(data)

#     def in_waiting(self):
#         n_waiting = len(self.fifo)
#         if not n_waiting:
#             data = self.tn.read_eager()
#             self.fifo.extend(data)
#             return len(data)
#         else:
#             return n_waiting


class Pyboard(object):
    serial: serial.Serial

    def __init__(self, device: str, baudrate: int = 115200):
        for _ in range(0, 3):
            try:
                self.serial = serial.Serial(device, baudrate=baudrate, interCharTimeout=1)
                break
            except:
                time.sleep(1)
        else:
            raise Exception(f'Failed to access {device}')

    def close(self):
        self.serial.close()

    def stdout_write_bytes(self, data: str):
        sys.stdout.buffer.write(data.replace(b'\x04', b''))
        sys.stdout.buffer.flush()

    def send_ok(self) -> bytes:
        self.serial.write(b'\x04')
        return b'OK'

    def send_ctrl_a(self) -> bytes:
        # ctrl-A: enter raw REPL
        self.serial.write(b'\x01')
        return b'raw REPL; CTRL-B to exit\r\n>'

    def send_ctrl_b(self):
        # ctrl-B: exit raw REPL
        self.serial.write(b'\x02')

    def send_ctrl_c(self, max_times: int = 2) -> bytes:
        # ctrl-C twice: interrupt any running program
        for _ in range(0, max_times):
            self.serial.write(b'\x03')
            time.sleep(0.1)

        return b'raw REPL; CTRL-B to exit\r\n'

    def send_ctrl_d(self) -> bytes:
        # ctrl-D: soft reset
        self.serial.write(b'\x04')
        return b'soft reboot\r\n'

    def read_until(self, delimiter: bytes, stream_output: bool = False) -> Optional[bytes]:
        data = self.serial.read(1)

        if stream_output:
            self.stdout_write_bytes(data)

        timeout = 0
        max_len = len(delimiter)

        while not timeout >= 1000:
            if data.endswith(delimiter):
                return data
            elif self.serial.inWaiting() > 0:
                timeout = 0
                stream_data = self.serial.read(1)
                data += stream_data

                if stream_output:
                    self.stdout_write_bytes(stream_data)
                    data = data[-max_len:]
            else:
                timeout += 1
                time.sleep(0.001)

    def __enter__(self):
        self.send_ctrl_c()
        n = self.serial.inWaiting()

        while n > 0:
            self.serial.read(n)
            n = self.serial.inWaiting()

        for _ in range(0, 5): 
            if self.read_until(self.send_ctrl_a()):
                break

            time.sleep(0.01)
        else:
            raise Exception('REPL: could not enter')

        if not self.read_until(self.send_ctrl_d()):
            raise Exception('REPL: could not soft reboot')

        if not self.read_until(self.send_ctrl_c()):
            raise Exception('REPL: could not interrupt after soft reboot')

        return self.__terminal

    def __exit__(self, a, b, c):
        self.send_ctrl_b()

    def __terminal(self, command: str, stream_output: bool = False, catch_output: bool = True) -> Optional[str]:
        command = textwrap.dedent(command)
        # send input
        if not isinstance(command, bytes):
            command = bytes(command, encoding='utf8')

        if not self.read_until(b'>'):
            raise Exception('Terminal: prompt has been lost')

        for i in range(0, len(command), BUFFER_SIZE):
            self.serial.write(command[i: min(i + BUFFER_SIZE, len(command))])
            time.sleep(0.001)

        if not self.read_until(self.send_ok()):
            raise Exception('Terminal: could not execute command')

        if not catch_output:
            return

        # catch output
        data = self.read_until(b'\x04', stream_output=stream_output)

        if not data:
            raise Exception('Terminal: timeout waiting for first EOF reception')

        exception = self.read_until(b'\x04')

        if not exception:
            raise Exception('Terminal: timeout waiting for second EOF reception')

        data, exception = (data[:-1].decode('utf-8'), exception[:-1].decode('utf-8'))

        if exception:
            reason = 'Traceback (most recent call last):'
            traceback = exception.split(reason)
            
            if len(traceback) == 2:
                for call in traceback[1][:-2].split('\\r\\n'):
                    reason += f'{call}\n'

                raise Exception(f'-------\n{reason.strip()}')
            else:
                raise Exception(f'-------\n{exception}')

        return data.strip()


class FileSystem(object):
    _pyboard: Pyboard

    def __init__(self, pyboard: Pyboard):
        self._pyboard = pyboard

    def checksum(self, source: str, data: str) -> str:
        output = self.terminal(f"""
            def checksum(data):
                v = 21
                for c in data.decode("utf-8"):
                    v ^= ord(c)
                return v
            
            with open("{source}", "rb") as fh:
                print(checksum(fh.read()))
        """)

        if isinstance(data, bytes):
            data = data.decode('utf-8')

        v = 21
        for c in data:
            v ^= ord(c)

        return f'{v}:{output}'

    def get(self, filename: str) -> Tuple[bytes, bool]:
        if not filename.startswith('/'):
            filename = '/' + filename

        output = self.terminal(f"""
            import sys
            import ubinascii
            with open('{filename}', 'rb') as infile:
                while True:
                    result = infile.read({BUFFER_SIZE})
                    if result == b'':
                        break
                    sys.stdout.write(ubinascii.hexlify(result))
        """)
        output = binascii.unhexlify(output)
        
        return (output, self.checksum(filename, output))

    def download(self, source: str, destination: str) -> str:
        output, checksum = self.get(source)

        with open(destination, 'wb') as fh:
            fh.write(output)
            
        return checksum

    def put(self, filename: str, data: bytes) -> Tuple[str, bool]:
        if not filename.startswith('/'):
            filename = '/' + filename

        if not isinstance(data, bytes):
            data = bytes(data, encoding='utf8')

        try:
            if os.path.dirname(filename):
                self.mkdir(os.path.dirname(filename))

            with self._pyboard as terminal:
                size = len(data)
                terminal(f"""fh = open("{filename}", "wb")""")

                for i in range(0, size, BUFFER_SIZE):
                    chunk_size = min(BUFFER_SIZE, size - i)
                    chunk = repr(data[i : i + chunk_size])
                    terminal(f"""fh.write({chunk})""")

                terminal("""fh.close()""")
        except Exception as e:
            raise e

        return (filename, self.checksum(filename, data))

    def upload(self, source: str, destination: str) -> str:
        with open(source, 'rb') as fh:
            _, checksum = self.put(destination, fh.read())

        return checksum

    def ls(self, dirname: str = '/') -> Tuple[int, List, str]:
        if not dirname.startswith('/'):
            dirname = '/' + dirname

        output = self.terminal(f"""
            try:        
                import os
                import json
            except ImportError:
                import uos as os
                import ujson as json

            def ls(dirname):
                e = []
                s = os.stat(dirname)
                if s[0] == 0x4000:
                    if not dirname.endswith("/"):
                        dirname += "/"
                    for t in os.ilistdir(dirname):
                        if t[1] == 0x4000:
                            e.append((dirname + t[0] + '/', -1))
                            e.extend(ls(dirname + t[0] + '/'))
                        else:
                            e.append((dirname + t[0], os.stat(dirname + t[0])[6]))
                else:
                    e.append((dirname, s[6]))
                return e
            
            try:
                s = 1
                r = ls("{dirname}")
                x = ''
            except Exception as e:
                s = 0
                r = [("{dirname}", -2)]
                x = str(e)

            print(json.dumps([s, r, x]))
        """)
        
        return json.loads(output)

    def rm(self, filename: str) -> Tuple[int, List, str]:
        if not filename.startswith('/'):
            filename = '/' + filename

        output = self.terminal(f"""
            try:        
                import os
                import json
            except ImportError:
                import uos as os
                import ujson as json

            def ls(dirname):
                e = []
                if not dirname.endswith("/"):
                    dirname += "/"
                for t in os.ilistdir(dirname):
                    if t[1] == 0x4000:
                        e.append((dirname + t[0] + '/', -1))
                        e.extend(ls(dirname + t[0] + '/'))
                    else:
                        e.append((dirname + t[0], os.stat(dirname + t[0])[6]))
                return e

            def rm(filename):
                r = []
                if os.stat(filename)[0] == 0x4000:
                    e = ls(filename)
                    e.append((filename, -1))
                    for f, s in e:
                        if not s == -1:
                            try:
                                os.remove(f)
                                r.append((f, 1, ''))
                            except Exception as e:
                                r.append((f, 0, str(e)))
                    for f, s in e:
                        if s == -1:
                            try:
                                os.rmdir(f)
                                r.append((f, 1, ''))
                            except Exception as e:
                                r.append((f, 0, str(e)))
                else:
                    try:
                        os.remove(filename)
                        r.append((filename, 1, ''))
                    except Exception as e:
                        r.append((filename, 0, str(e)))
                return r

            try:
                s = 1
                r = rm("{filename}")
                x = ''
            except Exception as e:
                s = 0
                r = [("{filename}", 0, str(e))]
                x = str(e)

            print(json.dumps([s, r, x]))
        """)

        return json.loads(output)

    def mkdir(self, dirname: str) -> List:
        if not dirname.startswith('/'):
            dirname = '/' + dirname
        
        if dirname.endswith('/'):
            dirname = dirname[:-1]

        output = self.terminal(f"""
            try:        
                import os
                import json
            except ImportError:
                import uos as os
                import ujson as json

            r = []
            d = []
            for zd in str("{dirname}").split("/"):
                if not zd:
                    continue
                d.append(zd)
                zd = "/".join(d)

                try:
                    os.mkdir(zd)
                    r.append(("/" + zd, 1))
                except Exception as e:
                    if str(e).find('EEXIST'):
                        r.append(("/" + zd, 1))
                    else:
                        r.append(("/" + zd, 0, str(e)))

            print(json.dumps(r))
        """)

        return json.loads(output)

    def launch(self, filename: str):
        try:
            self.terminal(f"""
                with open("{filename}", "r") as fh:
                    exec(fh.read())
            """, stream_output=True)
        except Exception as e:
            print(str(e))

    def terminal(self, command: str, stream_output: bool = False) -> str:
        with self._pyboard as terminal:
            try:
                return terminal(command, stream_output=stream_output)
            except Exception as e:
                raise e


PRINT_PREFIX_EXCEPTION = '\t =>'


class Picowatch(object):
    _fs: FileSystem

    def __init__(self, pyboard: Pyboard):
        self._fs = FileSystem(pyboard)

    def listing(self, remote: str = '/'):
        status, output, exception = self._fs.ls(remote)

        if status:
            for name, size in output:
                if size == -1:
                    print('[/]', name[1:])
                else:
                    print('[·]', name[1:], f'({size}b)')
        else:
            print('[?]', remote, PRINT_PREFIX_EXCEPTION, exception)

    def upload(self, local: str, remote: str = ''):
        pass

    def download(self, filepath: str):
        status, output, exception = self._fs.ls(filepath)

        if status:
            for remote, size in output:
                if size == -1:
                    os.makedirs(remote, 777, exist_ok=True)

            for remote, size in output:
                if not size == -1:
                    try:
                        print('[↓]', remote, PRINT_PREFIX_EXCEPTION, self._fs.download(remote, f'.{remote}'))
                    except Exception as e:
                        print('[?]', remote, PRINT_PREFIX_EXCEPTION, str(e))
        else:
            print('[?]', filepath, PRINT_PREFIX_EXCEPTION, exception)

    def delete(self, filepath: str):
        status, output, exception = self._fs.rm(filepath)

        if status:
            for remote, checked, message in output:
                if checked:
                    print('[-]', remote)
                else:
                    print('[?]', remote, PRINT_PREFIX_EXCEPTION, message)
        else:
            print('[?]', filepath, PRINT_PREFIX_EXCEPTION, exception)
       
    def execute(self, filepath: str):
        pass

    def watch(self, local: str):
        pass


# picowatch = Picowatch(Pyboard('COM5'))

# picowatch.listing('/lib/')