Make bytetree.py vbt-compatible

README.md

@@ -135,12 +135,21 @@ the results to <https://vmn-webapp.azurewebsites.net/results> (by default). Usag
 
 ```text
 usage: demo.py tally [-h] [--vote-collecting-server SERVER]
+                     [--use-bytetree-parser | --use-vbt | --skip-plaintexts]
 
 optional arguments:
   -h, --help            show this help message and exit
   --vote-collecting-server SERVER
                         Address of vote collecting server where the script
                         POSTs the public key and GETs the ciphertexts
+  --use-bytetree-parser
+                        Use bytetree.py to parse plaintexts. File must be
+                        located in directory '../webdemo/' relative to this
+                        script's location
+  --use-vbt             Use `vbt` to parse plaintexts. Must be available in
+                        $PATH
+  --skip-plaintexts     Do not parse the plaintexts and do not upload them to
+                        the results page
 ```
 
 ### Shutting down all servers

scripts/demo.py

@@ -36,6 +36,8 @@ def print_color(color, *args, **kwargs):
         "Python library `requests` is not installed, commands `start` and `tally` will not work."
     )
 
+BYTETREE_PATH = "../webdemo/"
+
 
 class VirtualMachine:
     def __init__(self, vm_id, args):
@@ -214,6 +216,26 @@ def parse_args():
         dest="server",
     )
 
+    # Tally
+    g = tally_election_parser.add_mutually_exclusive_group()
+    g.add_argument(
+        "--use-bytetree-parser",
+        action="store_true",
+        help=f"Use bytetree.py to parse plaintexts. File must be located in directory '{BYTETREE_PATH}' relative to this script's location",
+        dest="bytetree",
+    )
+    g.add_argument(
+        "--use-vbt",
+        action="store_true",
+        dest="vbt",
+        help="Use `vbt` to parse plaintexts. Must be available in $PATH",
+    )
+    g.add_argument(
+        "--skip-plaintexts",
+        action="store_true",
+        help="Do not parse the plaintexts and do not upload them to the results page",
+    )
+
     # Stop
     stop_parser.add_argument(
         "--delete",
@@ -316,6 +338,13 @@ def start_main(args):
 def tally_main(args):
     require_requests()
 
+    vbt_call = None
+    if args.bytetree:
+        vbt_call = import_bytetree()
+    elif not args.skip_plaintexts:
+        args.use_vbt = True
+        vbt_call = _check_output_vbt
+
     with open("ciphertexts", "wb") as f:
         r = requests.get(urljoin(args.server, "ciphertexts"))
         r.raise_for_status()
@@ -343,10 +372,11 @@ def tally_main(args):
     )
     assert os.path.exists("plaintexts")
 
-    vbt_json = vbt("plaintexts")
-    print(vbt_json)
-    r = requests.post(urljoin(args.server, "results"), json=vbt_json)
-    r.raise_for_status()
+    if vbt_call is not None:
+        vbt_json = vbt_count("plaintexts", vbt_call)
+        print(vbt_json)
+        r = requests.post(urljoin(args.server, "results"), json=vbt_json)
+        r.raise_for_status()
 
     # Verify
     for vm in vms:
@@ -374,6 +404,22 @@ def require_requests():
         sys.exit(1)
 
 
+def import_bytetree():
+    path = os.path.join(os.path.dirname(os.path.realpath(__file__)), BYTETREE_PATH)
+    try:
+        sys.path.append(path)
+        from bytetree import byte_array_byte_tree_to_json
+
+        def vbt_call(fname):
+            with open(fname, "rb") as f:
+                return json.loads(byte_array_byte_tree_to_json(bytearray(f.read())))
+
+        return vbt_call
+    except ImportError as e:
+        error(f"Could not load bytetree.py that should have been located in {path}")
+        raise e
+
+
 def get_vms(args, start=True):
     vms = azure_vms_by_tag(args)
     if not vms:
@@ -663,7 +709,7 @@ def azure_install_server(vm):
     return p.communicate(input=INSTALL_SCRIPT)
 
 
-def vbt(fname):
+def vbt_count(fname, vbt_call):
     valid_chars = " -_.,()" + string.ascii_letters + string.digits
 
     return Counter(
@@ -676,7 +722,7 @@ def vbt(fname):
                 for c in map(chr, bytes.fromhex(x[0]))
                 if c in valid_chars
             ),
-            _check_output_vbt(fname),
+            vbt_call(fname),
         )
     )
 

webdemo/bytetree.py

@@ -1,7 +1,11 @@
+#!/usr/bin/env python
 from collections.abc import Sequence
+from typing import ByteString, Iterator, List, Tuple, Union
 
 BYTEORDER = "big"
 
+ByteTreeValue = Union[ByteString, Sequence["ByteTree"]]
+
 
 class ByteTree:
     """
@@ -14,28 +18,32 @@ class ByteTree:
     NODE = 0
     LEAF = 1
 
-    def __init__(self, value):
+    def __init__(self, value: ByteTreeValue) -> None:
         if not isinstance(value, Sequence):
             raise TypeError("value should be of type Sequence")
 
+        self.type: int
         if isinstance(value[0], int):
             self.type = ByteTree.LEAF
         else:
             self.type = ByteTree.NODE
-        self.value = value
+        self.value: ByteTreeValue = value
 
-    def is_node(self):
+    def is_node(self) -> bool:
         return self.type == ByteTree.NODE
 
-    def is_leaf(self):
+    def is_leaf(self) -> bool:
         return self.type == ByteTree.LEAF
 
     @classmethod
-    def from_byte_array(cls, source):
+    def from_byte_array(cls, source: ByteString) -> "ByteTree":
+        """
+        Read a byte tree from a byte array
+        """
         return cls._from_byte_array(source, 0)[0]
 
     @classmethod
-    def _from_byte_array(cls, source, index=0):
+    def _from_byte_array(cls, source: ByteString, index=0) -> Tuple["ByteTree", int]:
         original_index = index
         tpe = source[index]
         assert tpe in (cls.NODE, cls.LEAF)
@@ -59,7 +67,10 @@ def _from_byte_array(cls, source, index=0):
             byte_tree = cls(children)
         return byte_tree, index - original_index
 
-    def to_byte_array(self):
+    def to_byte_array(self) -> ByteString:
+        """
+        Convert byte tree to its continuous byte array representation.
+        """
         byte_array = self.type.to_bytes(1, BYTEORDER)
         byte_array += len(self.value).to_bytes(4, BYTEORDER)
         # index = 0 + 1 + 4
@@ -69,8 +80,68 @@ def to_byte_array(self):
             # index += len(self.value)
         else:
             for child in self.value:
-                # child_bytes, offset = child.to_byte_array()
+                child: "ByteTree"
+
                 byte_array += child.to_byte_array()
                 # index += offset
 
         return byte_array
+
+    def pretty_str(self, indent: int = 0) -> str:
+        """
+        Writes formatted string illustrating the byte tree.
+
+        See `ByteTreeBasic::prettyWriteTo` in verificatum-vcr.
+        """
+        return "".join(self._pretty_str(indent))
+
+    def _pretty_str(self, indent: int) -> Iterator[str]:
+        s = 2 * indent * " "
+        if self.is_leaf():
+            data = "".join(_byte_to_hex(x) for x in self.value)
+            yield f'{s}"{data}"'
+        else:
+            yield f"{s}[\n"
+            for idx, child in enumerate(self.value):
+                child: "ByteTree"
+
+                is_last = idx == len(self.value) - 1
+
+                yield from child.pretty_str(indent + 1)
+                if not is_last:
+                    yield ","
+                yield "\n"
+            yield f"{s}]"
+
+
+def _byte_to_hex(x: int) -> str:
+    x = hex(x)[2:]
+    assert len(x) in (1, 2)
+    return x.rjust(2, "0")
+
+
+def _main(args: List[str]) -> int:
+    """
+    Port of `vbt`'s base functionality in python.
+    """
+    if len(args) > 2:
+        print("Usage:", args[0], "<filename>", file=sys.stderr)
+        return 1
+    elif len(args) == 2:
+        with open(args[1], "rb") as f:
+            inp = f.read()
+    else:
+        inp = sys.stdin.read()
+
+    print(byte_array_byte_tree_to_json(bytearray(inp)))
+    return 0
+
+
+def byte_array_byte_tree_to_json(ba: ByteString):
+    return ByteTree.from_byte_array(ba).pretty_str()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(_main(sys.argv))