SHARE
TWEET

Untitled

a guest Oct 13th, 2019 64 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #!/usr/bin/env python3
  2. import abc
  3. import argparse
  4. import contextlib
  5. import io
  6. import os
  7. import platform
  8. import socket
  9. import subprocess
  10. import sys
  11. import zlib
  12.  
  13. from typing import Any, Callable, Optional, Sequence, Tuple, Union, List
  14.  
  15.  
  16. class ArticulatedValueError(ValueError, argparse.ArgumentTypeError):
  17.     """ArticulatedTypeError is a ValueError that contains an error message
  18.     displayed by argparse."""
  19.  
  20.  
  21. def parse_host_and_maybe_port(host_port: str) -> Tuple[str, Optional[int]]:
  22.     """
  23.     parse_host_and_maybe_port splits host:port into a typed tuple.
  24.  
  25.     The tuple always contains an address but the port number may be None.  A
  26.     ValueError is raised if the port is provided but does not represent a
  27.     valid port number.
  28.     """
  29.     parts = host_port.split(":", 1)
  30.     host = parts[0]
  31.     port = None  # type: Optional[int]
  32.     if len(parts) == 2:
  33.         try:
  34.             port = int(parts[1])
  35.         except ValueError:
  36.             raise ArticulatedValueError("port is not a number")
  37.         if 0 > port > 65535:
  38.             raise ArticulatedValueError("port not in range 0..65535")
  39.     return (host, port)
  40.  
  41.  
  42. def parse_host_and_port(host_port: str) -> Tuple[str, int]:
  43.     """
  44.     parse_host_and_port splits host:port into a typed tuple.
  45.  
  46.     A ValueError is raised if the format is invalid.
  47.     """
  48.     host, maybe_port = parse_host_and_maybe_port(host_port)
  49.     if maybe_port is None:
  50.         # See note in parse_host_and_maybe_port
  51.         raise ArticulatedValueError("port number must be explicitly provided")
  52.     return (host, maybe_port)
  53.  
  54.  
  55. class Device(int):
  56.     """
  57.     Device is a device number with major and minor components.
  58.  
  59.     Note that this class does not attempt to mimic peculiar
  60.     encoding used by the Linux kernel.
  61.     """
  62.  
  63.     @classmethod
  64.     def pack(cls, major: int, minor: int) -> "Device":
  65.         return cls((major << 16) | (minor & (1 << 16) - 1))
  66.  
  67.     def __str__(self) -> str:
  68.         return "{}:{}".format(self.major, self.minor)
  69.  
  70.     def __repr__(self) -> str:
  71.         return "Device.pack({}, {})".format(self.major, self.minor)
  72.  
  73.     @property
  74.     def major(self) -> int:
  75.         """major is the higher 16 bits of the device number."""
  76.         return self >> 16
  77.  
  78.     @property
  79.     def minor(self) -> int:
  80.         """minor is the lower 16 bits of the device number."""
  81.         return self & ((1 << 16) - 1)
  82.  
  83.  
  84. class OptionalFields(List[str]):
  85.     def __str__(self) -> str:
  86.         """__str__ returns the special formatting of optional fields."""
  87.         if len(self):
  88.             return " ".join(self) + " -"
  89.         else:
  90.             return "-"
  91.  
  92.  
  93. class MountInfoEntry(object):
  94.     """Single entry in /proc/pid/mointinfo, see proc(5)"""
  95.  
  96.     mount_id: int
  97.     parent_id: int
  98.     dev: Device
  99.     root_dir: str
  100.     mount_point: str
  101.     mount_opts: str
  102.     opt_fields: OptionalFields
  103.     fs_type: str
  104.     mount_source: str
  105.     sb_opts: str
  106.  
  107.     def __init__(self) -> None:
  108.         self.mount_id = 0
  109.         self.parent_id = 0
  110.         self.dev = Device.pack(0, 0)
  111.         self.root_dir = ""
  112.         self.mount_point = ""
  113.         self.mount_opts = ""
  114.         self.opt_fields = OptionalFields()
  115.         self.fs_type = ""
  116.         self.mount_source = ""
  117.         self.sb_opts = ""
  118.  
  119.     def __eq__(self, other: object) -> "Union[NotImplemented, bool]":
  120.         if not isinstance(other, MountInfoEntry):
  121.             return NotImplemented
  122.         return (
  123.             self.mount_id == other.mount_id
  124.             and self.parent_id == other.parent_id
  125.             and self.dev == other.dev
  126.             and self.root_dir == other.root_dir
  127.             and self.mount_point == other.mount_point
  128.             and self.mount_opts == other.mount_opts
  129.             and self.opt_fields == other.opt_fields
  130.             and self.fs_type == other.fs_type
  131.             and self.mount_source == other.mount_source
  132.             and self.sb_opts == other.sb_opts
  133.         )
  134.  
  135.     @classmethod
  136.     def parse(cls, line: str) -> "MountInfoEntry":
  137.         it = iter(line.split())
  138.         self = cls()
  139.         self.mount_id = int(next(it))
  140.         self.parent_id = int(next(it))
  141.         dev_maj, dev_min = map(int, next(it).split(":"))
  142.         self.dev = Device((dev_maj << 16) | dev_min)
  143.         self.root_dir = next(it)
  144.         self.mount_point = next(it)
  145.         self.mount_opts = next(it)
  146.         self.opt_fields = OptionalFields()
  147.         for opt_field in it:
  148.             if opt_field == "-":
  149.                 break
  150.             self.opt_fields.append(opt_field)
  151.         self.fs_type = next(it)
  152.         self.mount_source = next(it)
  153.         self.sb_opts = next(it)
  154.         try:
  155.             next(it)
  156.         except StopIteration:
  157.             pass
  158.         else:
  159.             raise ValueError("leftovers after parsing {!r}".format(line))
  160.         return self
  161.  
  162.     def __str__(self) -> str:
  163.         return (
  164.             "{0.mount_id} {0.parent_id} {0.dev} {0.root_dir}"
  165.             " {0.mount_point} {0.mount_opts} {0.opt_fields} {0.fs_type}"
  166.             " {0.mount_source} {0.sb_opts}"
  167.         ).format(self)
  168.  
  169.     def __repr__(self) -> str:
  170.         return "MountInfoEntry.parse({!r})".format(str(self))
  171.  
  172.     @property
  173.     def dev_maj(self) -> int:
  174.         return self.dev.major
  175.  
  176.     @property
  177.     def dev_min(self) -> int:
  178.         return self.dev.minor
  179.  
  180.  
  181. class ReMountReadOnly:
  182.     """ReMountReadOnly is a context manager remounting filesystem to read-only."""
  183.  
  184.     def __init__(
  185.         self, path: str, on_remount: Optional[Callable[[str, bool], None]] = None
  186.     ):
  187.         self.path = path
  188.         self.altered = False
  189.         self.on_remount = on_remount
  190.  
  191.     def __enter__(self) -> "Optional[ReMountReadOnly]":
  192.         retcode = subprocess.call(["mount", "-o", "remount,ro", self.path])
  193.         if retcode == 0:
  194.             self.altered = True
  195.             if self.on_remount is not None:
  196.                 self.on_remount(self.path, True)
  197.             return self
  198.         return None
  199.  
  200.     def __exit__(self, *exc_details: Any) -> None:
  201.         if not self.altered:
  202.             return
  203.         retcode = subprocess.call(["mount", "-o", "remount,rw", self.path])
  204.         if retcode == 0 and self.on_remount is not None:
  205.             self.on_remount(self.path, False)
  206.  
  207.  
  208. class Command(abc.ABC):
  209.     """Command is the interface for command line commands."""
  210.  
  211.     class Args(argparse.Namespace):
  212.         """Args describes the parsed arguments."""
  213.  
  214.     @abc.abstractproperty
  215.     def name(self) -> str:
  216.         """name of the command"""
  217.  
  218.     @abc.abstractproperty
  219.     def description(self) -> str:
  220.         """description of the command"""
  221.  
  222.     @abc.abstractmethod
  223.     def configure_parser(self, parser: argparse.ArgumentParser) -> None:
  224.         """configure_parser configures the parser specific to the command."""
  225.  
  226.     @abc.abstractmethod
  227.     def invoke(self, args: Args) -> None:
  228.         """invoke executes the command with the given arguments."""
  229.  
  230.     def say(self, msg: str) -> None:
  231.         """say prints stuff as the command"""
  232.         print("raspi-tool", msg)
  233.  
  234.  
  235. class SendImgCmd(Command):
  236.  
  237.     name = "send-image"
  238.     description = "send an image from this device to a remote machine"
  239.  
  240.     class Args(Command.Args):
  241.         device: str
  242.         addr: Tuple[str, int]
  243.  
  244.     def configure_parser(self, parser: argparse.ArgumentParser) -> None:
  245.         parser.add_argument(
  246.             "--device",
  247.             metavar="DEVICE",
  248.             default="/dev/mmcblk0",
  249.             help="File representing SD card",
  250.         )
  251.         parser.add_argument(
  252.             "addr",
  253.             metavar="HOST:PORT",
  254.             type=parse_host_and_port,
  255.             help="Destination (e.g. running raspberry-pi-tool recv-image)",
  256.         )
  257.  
  258.     def _on_remounted(self, path: str, is_ro: bool) -> None:
  259.         self.say("{} re-mounted {}".format(path, "ro" if is_ro else "rw"))
  260.  
  261.     def invoke(self, args: Args) -> None:
  262.         with contextlib.ExitStack() as stack:
  263.             try:
  264.                 sink = socket.create_connection(args.addr)
  265.             except IOError as exc:
  266.                 raise SystemExit("cannot connect to {}: {}".format(args.addr, exc))
  267.             else:
  268.                 stack.enter_context(sink)
  269.  
  270.             with open("/proc/self/mountinfo") as stream:
  271.                 mountinfo = [MountInfoEntry.parse(line) for line in stream]
  272.             for entry in mountinfo:
  273.                 if entry.mount_source.startswith(args.device):
  274.                     stack.enter_context(
  275.                         ReMountReadOnly(entry.mount_point, self._on_remounted)
  276.                     )
  277.             try:
  278.                 source = io.FileIO(args.device)
  279.             except IOError as exc:
  280.                 raise SystemExit("cannot open {}: {}".format(args.device, exc))
  281.             else:
  282.                 stack.enter_context(source)
  283.  
  284.             self.say("sending {} to {}".format(args.device, args.addr))
  285.             packer = zlib.compressobj(level=9)
  286.             buf = bytearray(4096)
  287.             while True:
  288.                 n = source.readinto(buf)
  289.                 if n == 0:
  290.                     break
  291.                 zbuf = packer.compress(buf[:n])
  292.                 sink.sendall(zbuf)
  293.                 del zbuf
  294.             zbuf = packer.flush(zlib.Z_FINISH)
  295.             sink.sendall(zbuf)
  296.             del zbuf
  297.  
  298.             self.say("sent!")
  299.  
  300.  
  301. class RecvImgCmd(Command):
  302.  
  303.     name = "recv-image"
  304.     description = "receive an image a remote image and save it locally"
  305.  
  306.     class Args(Command.Args):
  307.         img: str
  308.         addr: Optional[Tuple[str, int]]
  309.  
  310.     def configure_parser(self, parser: argparse.ArgumentParser) -> None:
  311.         parser.add_argument(
  312.             "--addr",
  313.             metavar="HOST:PORT",
  314.             type=parse_host_and_port,
  315.             help="Destination (e.g. running raspberry-pi-tool recv-image)",
  316.         )
  317.         parser.add_argument(
  318.             "--image",
  319.             dest="img",
  320.             metavar="IMAGE",
  321.             help="File representing SD card image",
  322.             default="pi.img",
  323.         )
  324.  
  325.     def invoke(self, args: Args) -> None:
  326.         with contextlib.ExitStack() as stack:
  327.             # Open the file we want to write to.
  328.             try:
  329.                 sink = stack.enter_context(io.FileIO(args.img, "w"))
  330.             except IOError as exc:
  331.                 raise SystemExit("cannot open {}: {}".format(args.img, exc))
  332.             # Open a socket and find a port to bind to.
  333.             sock = stack.enter_context(
  334.                 socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  335.             )
  336.             host = "0.0.0.0"
  337.             port = 12345
  338.             if args.addr is not None:
  339.                 host, port = args.addr
  340.                 search = 1
  341.             else:
  342.                 search = 10
  343.             for try_port in range(port, port + search):
  344.                 try:
  345.                     sock.bind((host, try_port))
  346.                 except OSError as exc:
  347.                     self.say("cannot bind to {}:{}: {}".format(host, try_port, exc))
  348.                 else:
  349.                     break
  350.             else:
  351.                 raise SystemExit("cannot bind to port, giving up")
  352.             # Listen for incoming connections.
  353.             sock.listen()
  354.             self.say("listening on {}:{}".format(host, port))
  355.             hostname = socket.gethostname()
  356.             hostaddr = socket.gethostbyname(hostname)
  357.             if hostaddr.startswith("127."):
  358.                 self.say("cannot determine useful address of this machine")
  359.             else:
  360.                 self.say("on the raspberry pi issue the following command")
  361.                 self.say("$ sudo ./raspi-tool.py send-image {}:{}".format(hostaddr, port))
  362.             self.say("waiting for connection...")
  363.             # Wait for client connection.
  364.             try:
  365.                 source, remote_addr = sock.accept()
  366.             except Exception as exc:
  367.                 raise SystemExit("cannot accept connection: {}".format(exc))
  368.  
  369.             self.say("saving image from {} to {}".format(remote_addr, args.img))
  370.             packer = zlib.decompressobj()
  371.             buf = bytearray(4096)
  372.             while True:
  373.                 n = source.recv_into(buf, 0)
  374.                 if n == 0:
  375.                     break
  376.                 zbuf = packer.decompress(buf[:n])
  377.                 sink.write(zbuf)
  378.                 del zbuf
  379.             zbuf = packer.flush(zlib.Z_FINISH)
  380.             sink.write(zbuf)
  381.             del zbuf
  382.             self.say("saved!")
  383.  
  384.  
  385. def _make_parser(commands: Sequence[Command]) -> argparse.ArgumentParser:
  386.     """_make_parser creates an argument parser with given subcommands."""
  387.     parser = argparse.ArgumentParser(
  388.         description="Utility for working with Raspberry Pi devices."
  389.     )
  390.     parser.add_argument("--version", action="version", version="0.1")
  391.     sub = parser.add_subparsers()
  392.     for cmd in commands:
  393.         cmd_parser = sub.add_parser(cmd.name, help=cmd.description)
  394.         cmd.configure_parser(cmd_parser)
  395.         cmd_parser.set_defaults(invoke=cmd.invoke)
  396.     return parser
  397.  
  398.  
  399. def main() -> None:
  400.     parser = _make_parser([SendImgCmd(), RecvImgCmd()])
  401.     # On windows, when invoked without arguments, default to receiving an image.
  402.     if platform.win32_ver()[0] != "" and len(sys.argv) == 1:
  403.         args = parser.parse_args(["recv-image"])
  404.     else:
  405.         args = parser.parse_args()
  406.     if not hasattr(args, "invoke"):
  407.         parser.error("select command to execute")
  408.     args.invoke(args)
  409.  
  410.  
  411. if __name__ == "__main__":
  412.     try:
  413.         main()
  414.     except KeyboardInterrupt:
  415.         pass
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top