more bacloud work

tools/bacloud

@@ -29,7 +29,7 @@ import sys
 import os
 from pathlib import Path
 from typing import TYPE_CHECKING
-from dataclasses import dataclass, asdict
+from dataclasses import dataclass
 import json
 import subprocess
 import tempfile
@@ -72,34 +72,36 @@ class Response:
       message: If present, client should print this message before any other
         response processing (including error handling) occurs.
       error: If present, client should abort with this error message.
-      loadpackage: If present, client should load this package from a
-        location on disk (arg1) and push its manifest to a server command
-        (arg2) with provided args (arg3). The manifest should be added to
-        the args as 'manifest'. arg4, if present, is the index file name whose
-        contents should be included with the manifest.
-      uploads: If present, client should upload the requested files (arg1)
-        from the loaded package to a server command (arg2) with provided
-        args (arg3). Arg4 and arg5 are a server call and args which should
-        be called once all file uploads finish.
       login: If present, a token that should be stored client-side and passed
         with subsequent commands.
       logout: If True, any existing client-side token should be discarded.
-      inline_downloads: If present, pathnames mapped to base64 gzipped data to
+      dirmanifest: If present, client should generate a manifest of this dir.
+        It should be added to endcommand args as 'manifest'.
+      uploads: If present, client should upload the requested files (arg1)
+        individually to a server command (arg2) with provided args (arg3).
+      uploads_inline: If present, a list of pathnames that should be base64
+        gzipped and uploaded to an 'uploads_inline' dict in endcommand args.
+        This should be limited to relatively small files.
+      downloads_inline: If present, pathnames mapped to base64 gzipped data to
         be written to the client. This should only be used for relatively
         small files as they are all included inline as part of the response.
       deletes: If present, file paths that should be deleted on the client.
       endmessage: If present, a message that should be printed after all other
         response processing is done.
+      endcommand: If present, this command is run with these args at the end
+        of response processing.
     """
     message: Optional[str] = None
     error: Optional[str] = None
-    loadpackage: Optional[Tuple[str, str, Dict, Optional[str]]] = None
-    uploads: Optional[Tuple[List[str], str, Dict, str, Dict]] = None
     login: Optional[str] = None
     logout: bool = False
-    inline_downloads: Optional[Dict[str, str]] = None
+    dirmanifest: Optional[str] = None
+    uploads: Optional[Tuple[List[str], str, Dict]] = None
+    uploads_inline: Optional[List[str]] = None
+    downloads_inline: Optional[Dict[str, str]] = None
     deletes: Optional[List[str]] = None
     endmessage: Optional[str] = None
+    endcommand: Optional[Tuple[str, Dict]] = None
 
 
 class CleanError(Exception):
@@ -175,8 +177,8 @@ class App:
 
     def __init__(self) -> None:
         self._state = StateData()
-        # self._package: Optional[Package] = None
         self._project_root: Optional[Path] = None
+        self._end_command_args: Dict = {}
 
     def run(self) -> None:
         """Run the tool."""
@@ -271,9 +273,7 @@ class App:
 
     def _upload_file(self, filename: str, call: str, args: Dict) -> None:
         print(f'{CLRBLU}Uploading {filename}{CLREND}', flush=True)
-        # assert self._package is not None
         with tempfile.TemporaryDirectory() as tempdir:
-            # srcpath = Path(self._package.path, filename)
             srcpath = Path(filename)
             gzpath = Path(tempdir, 'file.gz')
             subprocess.run(f'gzip --stdout "{srcpath}" > "{gzpath}"',
@@ -287,44 +287,48 @@ class App:
                     files=putfiles,
                 )
 
-    def _handle_loadpackage_response(
-            self, response: Response) -> Optional[Tuple[str, Dict]]:
-        assert response.loadpackage is not None
-        assert len(response.loadpackage) == 4
-        (packagepath, callname, callargs, indexfile) = response.loadpackage
-        assert isinstance(packagepath, str)
-        assert isinstance(callname, str)
-        assert isinstance(callargs, dict)
-        assert indexfile is None or isinstance(indexfile, str)
-        package = Package.load_from_disk(Path(packagepath))
+    def _handle_dirmanifest_response(self, response: Response) -> None:
+        from dataclasses import asdict
+        assert response.dirmanifest is not None
+        # assert len(response.dirmanifest) == 2
+        # (packagepath, indexfile) = response.dirmanifest
+        assert isinstance(response.dirmanifest, str)
+        # assert isinstance(callname, str)
+        # assert isinstance(callargs, dict)
+        # assert indexfile is None or isinstance(indexfile, str)
+        package = Package.load_from_disk(Path(response.dirmanifest))
 
         # Make the remote call they gave us with the package
         # manifest added in.
-        if indexfile is not None:
-            with Path(package.path, indexfile).open() as infile:
-                index = infile.read()
-        else:
-            index = ''
-        callargs['manifest'] = {
-            'index': index,
+        # if indexfile is not None:
+        #     with Path(package.path, indexfile).open() as infile:
+        #         index = infile.read()
+        # else:
+        #     index = ''
+        # callargs['manifest'] = {
+        #     'index': index,
+        #     'files': {key: asdict(val)
+        #               for key, val in package.files.items()}
+        # }
+
+        # Store the manifest to be included with our next called command.
+        self._end_command_args['manifest'] = {
             'files': {key: asdict(val)
                       for key, val in package.files.items()}
         }
-        return callname, callargs
+        # return callname, callargs
 
-    def _handle_uploads(self,
-                        response: Response) -> Optional[Tuple[str, Dict]]:
+    def _handle_uploads(self, response: Response) -> None:
         from concurrent.futures import ThreadPoolExecutor
         assert response.uploads is not None
-        # assert self._package is not None
-        assert len(response.uploads) == 5
-        (filenames, uploadcmd, uploadargs, completecmd,
-         completeargs) = response.uploads
+        assert len(response.uploads) == 3
+        filenames, uploadcmd, uploadargs = response.uploads
         assert isinstance(filenames, list)
         assert isinstance(uploadcmd, str)
         assert isinstance(uploadargs, dict)
-        assert isinstance(completecmd, str)
-        assert isinstance(completeargs, dict)
+
+        # assert isinstance(completecmd, str)
+        # assert isinstance(completeargs, dict)
 
         def _do_filename(filename: str) -> None:
             self._upload_file(filename, uploadcmd, uploadargs)
@@ -338,14 +342,14 @@ class App:
             list(executor.map(_do_filename, filenames))
 
         # Lastly, run the 'upload complete' command we were passed.
-        return completecmd, completeargs
+        # return completecmd, completeargs
 
-    def _handle_inline_downloads(self, response: Response) -> None:
+    def _handle_downloads_inline(self, response: Response) -> None:
         """Handle inline file data to be saved to the client."""
         import base64
         import zlib
-        assert response.inline_downloads is not None
-        for fname, fdata in response.inline_downloads.items():
+        assert response.downloads_inline is not None
+        for fname, fdata in response.downloads_inline.items():
             data_zipped = base64.b64decode(fdata)
             data = zlib.decompress(data_zipped)
             with open(fname, 'wb') as outfile:
@@ -357,6 +361,22 @@ class App:
         for fname in response.deletes:
             os.unlink(fname)
 
+    def _handle_uploads_inline(self, response: Response) -> None:
+        """Handle uploading files inline."""
+        import base64
+        import zlib
+        assert response.uploads_inline is not None
+        files: Dict[str, str] = {}
+        for filepath in response.uploads_inline:
+            if not os.path.exists(filepath):
+                raise CleanError(f'File not found: {filepath}')
+            with open(filepath, 'rb') as infile:
+                data = infile.read()
+            data_zipped = zlib.compress(data)
+            data_base64 = base64.b64encode(data_zipped).decode()
+            files[filepath] = data_base64
+        self._end_command_args['uploads_inline'] = files
+
     def run_user_command(self, args: List[str]) -> None:
         """Run a single user command to completion."""
 
@@ -364,24 +384,29 @@ class App:
 
         # Now talk to the server in a loop until there's nothing left to do.
         while nextcall is not None:
+            self._end_command_args = {}
             response = self._servercmd(*nextcall)
             nextcall = None
-            if response.loadpackage is not None:
-                nextcall = self._handle_loadpackage_response(response)
-            if response.uploads is not None:
-                nextcall = self._handle_uploads(response)
             if response.login is not None:
                 self._state.login_token = response.login
             if response.logout:
                 self._state.login_token = None
-            if response.inline_downloads:
-                self._handle_inline_downloads(response)
+            if response.dirmanifest is not None:
+                self._handle_dirmanifest_response(response)
+            if response.uploads_inline is not None:
+                self._handle_uploads_inline(response)
+            if response.uploads is not None:
+                self._handle_uploads(response)
+            if response.downloads_inline:
+                self._handle_downloads_inline(response)
             if response.deletes:
                 self._handle_deletes(response)
-
-            # This should always be printed last.
             if response.endmessage is not None:
                 print(response.endmessage, flush=True)
+            if response.endcommand is not None:
+                nextcall = response.endcommand
+                for key, val in self._end_command_args.items():
+                    nextcall[1][key] = val
 
 
 if __name__ == '__main__':