From dc92b484e88769cabd333bb84f224a3399694634 Mon Sep 17 00:00:00 2001
From: jpic <jpic@yourlabs.org>
Date: Mon, 3 Feb 2020 05:25:04 +0100
Subject: [PATCH] Asyncio rewrite

---
 podctl/__init__.py          |  2 +
 podctl/build.py             | 96 ++++++++++++++++++++++++++-----------
 podctl/console_script.py    | 20 ++------
 podctl/container.py         | 49 ++++++++++++++-----
 podctl/script.py            |  4 --
 podctl/visitable.py         | 16 +++++--
 podctl/visitors/base.py     |  8 ++++
 podctl/visitors/cmd.py      |  4 +-
 podctl/visitors/commit.py   |  8 ++--
 podctl/visitors/copy.py     | 10 ++--
 podctl/visitors/mount.py    |  4 +-
 podctl/visitors/npm.py      |  6 +--
 podctl/visitors/packages.py | 36 +++++++-------
 podctl/visitors/pip.py      | 33 ++++++-------
 podctl/visitors/run.py      |  4 +-
 podctl/visitors/template.py |  6 +--
 podctl/visitors/user.py     | 19 ++------
 17 files changed, 188 insertions(+), 137 deletions(-)

diff --git a/podctl/__init__.py b/podctl/__init__.py
index d54078a..e8f7ad1 100644
--- a/podctl/__init__.py
+++ b/podctl/__init__.py
@@ -1,3 +1,5 @@
+from .build import Build  # noqa
 from .container import Container  # noqa
 from .pod import Pod  # noqa
+from .script import Script  # noqa
 from .visitors import *  # noqa
diff --git a/podctl/build.py b/podctl/build.py
index 9331bcb..08cecf1 100644
--- a/podctl/build.py
+++ b/podctl/build.py
@@ -1,36 +1,71 @@
 import asyncio
 import os
+import asyncio
+import signal
 import subprocess
 import textwrap
 
 from .script import Script
 
 
+class WrongResult(Exception):
+    pass
+
+
 class Build(Script):
     def __init__(self, container):
         super().__init__()
         self.container = container
+        self.log = []
+        self.mounts = dict()
 
-    def append(self, value):
-        res = []
+    async def cmd(self, line):
+        log = dict(cmd=line)
+        self.log.append(log)
+        line = self.unshare(line)
+        print(self.container.name + ' | ' + line)
+        def protocol_factory():
+            from .console_script import BuildStreamProtocol
+            return BuildStreamProtocol(
+                self.container,
+                limit=asyncio.streams._DEFAULT_LIMIT,
+                loop=self.loop,
+            )
+        transport, protocol = await self.loop.subprocess_shell(
+            protocol_factory,
+            line,
+        )
+        proc = asyncio.subprocess.Process(
+            transport,
+            protocol,
+            self.loop,
+        )
+        result = await proc.wait()
+        if result:
+            raise WrongResult(proc)
+        return proc
+        '''
+        proc = await asyncio.create_subprocess_shell(
+            line,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        '''
+
+    async def append(self, value):
         for line in value.split('\n'):
             if line.startswith('#') or not line.strip():
                 continue
-            res.append(self.unshare(line))
-        return '\n'.join(res)
+            log = dict(proc=await self.cmd(line))
+            log['stdout'], log['stderr'] = await log['proc'].communicate()
 
-    async def unshare(self, line):
-        print('+ buildah unshare ' + line)
-        proc = await asyncio.create_subprocess_shell(
-            'buildah unshare ' + line,
-            stdout=asyncio.subprocess.PIPE,
-            stderr=asyncio.subprocess.PIPE,
-        ).decode('utf8')
-        stdout, stderr = await proc.communicate()
-        return stdout
+        return log['stdout']
 
-    def config(self, line):
-        self.append(f'buildah config {line} {self.ctr}')
+    def unshare(self, line):
+        return 'buildah unshare ' + line
+
+    async def config(self, line):
+        return await self.append(f'buildah config {line} {self.ctr}')
 
     def _run(self, cmd, inject=False):
         user = self.container.variable('username')
@@ -54,20 +89,25 @@ class Build(Script):
         else:
             return f'buildah run {self.ctr} -- {_cmd}'
 
-    def run(self, cmd):
-        self.append(self._run(cmd))
+    async def run(self, cmd):
+        return await self.append(self._run(cmd))
 
-    def copy(self, src, dst):
-        self.append(f'buildah copy {self.ctr} {src} {dst}')
+    async def copy(self, src, dst):
+        return await self.append(f'buildah copy {self.ctr} {src} {dst}')
 
-    def mount(self, src, dst):
-        self.run('sudo mkdir -p ' + dst)
-        self.append('mkdir -p ' + src)
-        self.append(f'mount -o bind {src} {self.mnt}{dst}')
-        #self.append('mounts=("$mnt' + dst + '" "${mounts[@]}")')
-        self.mounts.append((src, dst))
+    async def mount(self, src, dst):
+        await self.run('sudo mkdir -p ' + dst)
+        await self.append('mkdir -p ' + src)
+        await self.append(f'mount -o bind {src} {self.mnt}{dst}')
+        #await self.append('mounts=("$mnt' + dst + '" "${mounts[@]}")')
+        self.mounts[dst] = src
 
-    def umounts(self):
+    async def umounts(self):
         for src, dst in self.mounts:
-            self.append('buildah unmount ' + dst)
-        self.append('buildah unmount ' + self.ctr)
+            await self.append('buildah unmount ' + dst)
+        await self.append('buildah unmount ' + self.ctr)
+
+    def which(self, cmd):
+        for path in self.container.paths:
+            if os.path.exists(os.path.join(path, cmd)):
+                return True
diff --git a/podctl/console_script.py b/podctl/console_script.py
index f4ab7ad..3aa4793 100644
--- a/podctl/console_script.py
+++ b/podctl/console_script.py
@@ -42,24 +42,14 @@ async def build(*services, **kwargs):
     else:
         services = console_script.pod.services
 
+    a = []
     loop = asyncio.events.get_event_loop()
-    debug = console_script.parser.options.get('debug', False)
-    def protocol_factory():
-        return BuildStreamProtocol(
-            service,
-            limit=asyncio.streams._DEFAULT_LIMIT,
-            loop=loop,
-        )
-
     for name, service in services.items():
-        container = service.container
-        if not container.variable('base'):
-            continue
-        await container.build(loop, protocol_factory)
+        service.container.name = name
+        a.append(service.container.script('build', loop))
 
-    for proc in procs:
-        if proc.returncode != 0:
-            sys.exit(proc.returncode)
+    result = await asyncio.gather(*a, return_exceptions=False)
+    print(result)
 
 
 class ConsoleScript(cli2.ConsoleScript):
diff --git a/podctl/container.py b/podctl/container.py
index a3b4281..59420ce 100644
--- a/podctl/container.py
+++ b/podctl/container.py
@@ -1,19 +1,53 @@
+import asyncio
 import os
 
 from .build import Build
 from .visitable import Visitable
 
 
+class BuildStreamProtocol(asyncio.subprocess.SubprocessStreamProtocol):
+    def __init__(self, service, *args, **kwargs):
+        self.service = service
+        super().__init__(*args, **kwargs)
+
+    def pipe_data_received(self, fd, data):
+        if fd in (1, 2):
+            for line in data.split(b'\n'):
+                if not line:
+                    continue
+                sys.stdout.buffer.write(
+                    self.service.name.encode('utf8') + b' | ' + line + b'\n'
+                )
+            sys.stdout.flush()
+        super().pipe_data_received(fd, data)
+
+
+def protocol_factory():
+    loop = asyncio.events.get_event_loop()
+    return BuildStreamProtocol(
+        service,
+        limit=asyncio.streams._DEFAULT_LIMIT,
+        loop=loop,
+    )
+
+
 class Container(Visitable):
     default_scripts = dict(
         build=Build,
     )
+    paths = [
+        '/bin',
+        '/sbin',
+        '/usr/bin',
+        '/usr/sbin',
+    ]
 
-    def script(self, name):
+    async def script(self, name, loop):
         self.packages = []
         for visitor in self.visitors:
             self.packages += getattr(visitor, 'packages', [])
-        return super().script(name)
+        result = await super().script(name, loop)
+        return result
 
     def script_run(self, name, debug):
         script = f'.podctl_build_{name}.sh'
@@ -27,14 +61,3 @@ class Container(Visitable):
 
         x = 'x' if debug else ''
         return prefix + f'bash -eu{x} {script}'
-
-    async def build(self, loop, protocol_factory):
-        transport, protocol = await loop.subprocess_shell(
-            protocol_factory,
-            cmd,
-        )
-        await asyncio.subprocess.Process(
-            transport,
-            protocol,
-            loop,
-        ).communicate()
diff --git a/podctl/script.py b/podctl/script.py
index 9bc1ea4..297e63e 100644
--- a/podctl/script.py
+++ b/podctl/script.py
@@ -2,10 +2,6 @@ import textwrap
 
 
 class Script(list):
-    def __init__(self, shebang=None):
-        super().__init__()
-        self.append(shebang or '#/usr/bin/env bash')
-
     def __str__(self):
         if not getattr(self, '_postconfig', False):
             if hasattr(self, 'post_config'):
diff --git a/podctl/visitable.py b/podctl/visitable.py
index ffa668b..a62f77b 100644
--- a/podctl/visitable.py
+++ b/podctl/visitable.py
@@ -10,17 +10,23 @@ class Visitable:
             k: v(self) for k, v in self.default_scripts.items()
         }
 
-    def script(self, name):
+    async def script(self, name, loop):
         script = copy(self.scripts[name])
+        script.loop = loop
+        results = []
 
         for prefix in ('init_', 'pre_', '', 'post_'):
             method = prefix + name
             for visitor in self.visitors:
-                if hasattr(visitor, method):
-                    script.append(f'echo "{type(visitor).__name__}.{method}"')
-                    getattr(visitor, method)(script)
+                if not hasattr(visitor, method):
+                    continue
+
+                rep = {k: v if not isinstance(v, object) else type(v).__name__ for k, v in visitor.__dict__.items()}
+                print(self.name + ' | ', type(visitor).__name__, method, rep)
+                result = getattr(visitor, method)(script)
+                if result:
+                    await result
 
-        return script
 
     def visitor(self, name):
         for visitor in self.visitors:
diff --git a/podctl/visitors/base.py b/podctl/visitors/base.py
index 46c7aa9..1601328 100644
--- a/podctl/visitors/base.py
+++ b/podctl/visitors/base.py
@@ -3,3 +3,11 @@
 class Base:
     def __init__(self, base):
         self.base = base
+
+    async def init_build(self, script):
+        ctr = await script.cmd('buildah from ' + self.base)
+        stdout, stderr = await ctr.communicate()
+        script.ctr = stdout.decode('utf8').strip()
+        mnt = await script.cmd('buildah mount ' + script.ctr)
+        stdout, stderr = await mnt.communicate()
+        script.mnt = stdout.decode('utf8').strip()
diff --git a/podctl/visitors/cmd.py b/podctl/visitors/cmd.py
index 024662d..c488348 100644
--- a/podctl/visitors/cmd.py
+++ b/podctl/visitors/cmd.py
@@ -2,6 +2,6 @@ class Cmd:
     def __init__(self, cmd):
         self.cmd = cmd
 
-    def build(self, script):
+    async def build(self, script):
         # script._run() does not really support sudo code
-        script.run(self.cmd)
+        await script.run(self.cmd)
diff --git a/podctl/visitors/commit.py b/podctl/visitors/commit.py
index 37e24f5..3aa9f27 100644
--- a/podctl/visitors/commit.py
+++ b/podctl/visitors/commit.py
@@ -42,8 +42,8 @@ class Commit:
         # filter out tags which resolved to None
         self.tags = [t for t in self.tags if t is not None]
 
-    def post_build(self, script):
-        script.append(f'''
+    async def post_build(self, script):
+        await script.append(f'''
             umounts
             buildah commit --format={self.format} $ctr {self.repo}
         ''')
@@ -53,7 +53,7 @@ class Commit:
 
         if self.tags:
             tags = ' '.join([f'{self.repo}:{tag}' for tag in self.tags])
-            script.append(f'buildah tag {self.repo} {tags}')
+            await script.append(f'buildah tag {self.repo} {tags}')
 
             if self.push:
                 user = os.getenv('DOCKER_USER')
@@ -66,4 +66,4 @@ class Commit:
                     ])
 
                 for tag in self.tags:
-                    script.append(f'podman push {self.repo}:{tag}')
+                    await script.append(f'podman push {self.repo}:{tag}')
diff --git a/podctl/visitors/copy.py b/podctl/visitors/copy.py
index ebd2cfb..a8c51cf 100644
--- a/podctl/visitors/copy.py
+++ b/podctl/visitors/copy.py
@@ -13,13 +13,13 @@ class Copy:
             self.dst, self.mode = self.dst.split(':')
             self.owner = script.variable('user')
 
-    def build(self, script):
-        script.run(f'sudo mkdir -p {self.dst}')
+    async def build(self, script):
+        await script.run(f'sudo mkdir -p {self.dst}')
         for item in self.src:
-            script.append(f'cp -a {item} $mnt{self.dst}')
+            await script.append(f'cp -a {item} $mnt{self.dst}')
 
         if self.mode:
-            script.run(f'sudo chmod {self.mode} $mnt{self.dst}')
+            await script.run(f'sudo chmod {self.mode} $mnt{self.dst}')
 
         if self.owner:
-            script.run(f'sudo chown -R {self.owner} $mnt{self.dst}')
+            await script.run(f'sudo chown -R {self.owner} $mnt{self.dst}')
diff --git a/podctl/visitors/mount.py b/podctl/visitors/mount.py
index 1d904b9..1f09ac4 100644
--- a/podctl/visitors/mount.py
+++ b/podctl/visitors/mount.py
@@ -3,5 +3,5 @@ class Mount:
         self.src = src
         self.dst = dst
 
-    def build(self, script):
-        script.mount(self.src, self.dst)
+    async def build(self, script):
+        await script.mount(self.src, self.dst)
diff --git a/podctl/visitors/npm.py b/podctl/visitors/npm.py
index e8d6141..78ce9d1 100644
--- a/podctl/visitors/npm.py
+++ b/podctl/visitors/npm.py
@@ -2,9 +2,9 @@ class Npm:
     def __init__(self, install=None):
         self.npm_install = install
 
-    def build(self, script):
-        script.run('sudo npm update -g npm')
-        script.run(f'''
+    async def build(self, script):
+        await script.run('sudo npm update -g npm')
+        await script.run(f'''
         cd {self.npm_install}
         npm install
         ''')
diff --git a/podctl/visitors/packages.py b/podctl/visitors/packages.py
index 60ae617..522a5b9 100644
--- a/podctl/visitors/packages.py
+++ b/podctl/visitors/packages.py
@@ -37,7 +37,7 @@ class Packages:
         else:
             self.cache = os.path.join(os.getenv('HOME'), '.cache', self.mgr)
 
-    def pre_build(self, script):
+    async def pre_build(self, script):
         base = script.container.variable('base')
         if self.mgr:
             self.cmds = self.mgrs[self.mgr]
@@ -54,23 +54,23 @@ class Packages:
         if not self.mgr:
             raise Exception('Packages does not yet support this distro')
 
-    def build(self, script):
+    async def build(self, script):
         if not getattr(script.container, '_packages_upgraded', None):
             # run pkgmgr_setup functions ie. apk_setup
-            getattr(self, self.mgr + '_setup')(script)
+            await getattr(self, self.mgr + '_setup')(script)
             # first run on container means inject visitor packages
             self.packages += script.container.packages
-            script.run(self.cmds['upgrade'])
+            await script.run(self.cmds['upgrade'])
             script.container._packages_upgraded = True
 
-        script.run(' '.join([self.cmds['install']] + self.packages))
+        await script.run(' '.join([self.cmds['install']] + self.packages))
 
-    def apk_setup(self, script):
-        script.mount(self.cache, f'/var/cache/{self.mgr}')
+    async def apk_setup(self, script):
+        await script.mount(self.cache, f'/var/cache/{self.mgr}')
         # special step to enable apk cache
-        script.run('ln -s /var/cache/apk /etc/apk/cache')
-        script.append(f'''
-        old="$(find .cache/apk/ -name APKINDEX.* -mtime +3)"
+        await script.run('ln -s /var/cache/apk /etc/apk/cache')
+        await script.append(f'''
+        old="$(find {self.cache} -name APKINDEX.* -mtime +3)"
         if [ -n "$old" ] || ! ls .cache/apk/APKINDEX.*; then
             {script._run(self.cmds['update'])}
         else
@@ -78,18 +78,18 @@ class Packages:
         fi
         ''')
 
-    def dnf_setup(self, script):
-        script.mount(self.cache, f'/var/cache/{self.mgr}')
-        script.run('sh -c "echo keepcache=True >> /etc/dnf/dnf.conf"')
+    async def dnf_setup(self, script):
+        await script.mount(self.cache, f'/var/cache/{self.mgr}')
+        await script.run('sh -c "echo keepcache=True >> /etc/dnf/dnf.conf"')
 
-    def apt_setup(self, script):
+    async def apt_setup(self, script):
         cache = self.cache + '/$(source $mnt/etc/os-release; echo $VERSION_CODENAME)/'  # noqa
-        script.run('sudo rm /etc/apt/apt.conf.d/docker-clean')
+        await script.run('sudo rm /etc/apt/apt.conf.d/docker-clean')
         cache_archives = os.path.join(self.cache, 'archives')
-        script.mount(cache_archives, f'/var/cache/apt/archives')
+        await script.mount(cache_archives, f'/var/cache/apt/archives')
         cache_lists = os.path.join(self.cache, 'lists')
-        script.mount(cache_lists, f'/var/lib/apt/lists')
-        script.append(f'''
+        await script.mount(cache_lists, f'/var/lib/apt/lists')
+        await script.append(f'''
         old="$(find {cache_lists} -name lastup -mtime +3)"
         if [ -n "$old" ] || ! ls {cache_lists}/lastup; then
             until [ -z $(lsof /var/lib/dpkg/lock) ]; do sleep 1; done
diff --git a/podctl/visitors/pip.py b/podctl/visitors/pip.py
index 32b6fac..44eaa81 100644
--- a/podctl/visitors/pip.py
+++ b/podctl/visitors/pip.py
@@ -1,3 +1,4 @@
+from glob import glob
 import os
 
 
@@ -7,34 +8,28 @@ class Pip:
         self.pip = pip
         self.requirements = requirements
 
-    def build(self, script):
-        if self.pip:
-            script.append('_pip=' + self.pip)
-        else:
-            script.append(f'''
-            if {script._run("bash -c 'type pip3'")}; then
-                _pip=pip3
-            elif {script._run("bash -c 'type pip'")}; then
-                _pip=pip
-            elif {script._run("bash -c 'type pip2'")}; then
-                _pip=pip2
-            fi
-            ''')
+    async def build(self, script):
+        for pip in ('pip3', 'pip', 'pip2'):
+            if script.which(pip):
+                self.pip = pip
+                break
+
         if 'CACHE_DIR' in os.environ:
             cache = os.path.join(os.getenv('CACHE_DIR'), 'pip')
         else:
             cache = os.path.join(os.getenv('HOME'), '.cache', 'pip')
-        script.mount(cache, '/root/.cache/pip')
-        script.run('sudo $_pip install --upgrade pip')
+
+        await script.mount(cache, '/root/.cache/pip')
+        await script.run(f'sudo {self.pip} install --upgrade pip')
         source = [p for p in self.pip_packages if p.startswith('/')]
         if source:
-            script.run(
-                f'sudo $_pip install --upgrade --editable {" ".join(source)}'
+            await script.run(
+                f'sudo {self.pip} install --upgrade --editable {" ".join(source)}'
             )
 
         nonsource = [p for p in self.pip_packages if not p.startswith('/')]
         if nonsource:
-            script.run(f'sudo $_pip install --upgrade {" ".join(source)}')
+            await script.run(f'sudo {self.pip} install --upgrade {" ".join(source)}')
 
         if self.requirements:
-            script.run(f'sudo $_pip install --upgrade -r {self.requirements}')
+            await script.run(f'sudo {self.pip} install --upgrade -r {self.requirements}')
diff --git a/podctl/visitors/run.py b/podctl/visitors/run.py
index 6e16577..0da7e3d 100644
--- a/podctl/visitors/run.py
+++ b/podctl/visitors/run.py
@@ -2,6 +2,6 @@ class Run:
     def __init__(self, *commands):
         self.commands = commands
 
-    def build(self, script):
+    async def build(self, script):
         for command in self.commands:
-            script.run(command)
+            await script.run(command)
diff --git a/podctl/visitors/template.py b/podctl/visitors/template.py
index 00f702d..971f2f2 100644
--- a/podctl/visitors/template.py
+++ b/podctl/visitors/template.py
@@ -11,10 +11,10 @@ class Template:
         self.lines = lines
         self.variables = variables
 
-    def build(self, script):
+    async def build(self, script):
         self.script = '\n'.join([
             dedent(l).strip() for l in self.lines
         ]).format(**self.variables)
-        script.run(CMD.strip().format(**self.__dict__))
+        await script.run(CMD.strip().format(**self.__dict__))
         if self.script.startswith('#!'):
-            script.run('sudo chmod +x ' + self.target)
+            await script.run('sudo chmod +x ' + self.target)
diff --git a/podctl/visitors/user.py b/podctl/visitors/user.py
index 7b99f46..a5b4993 100644
--- a/podctl/visitors/user.py
+++ b/podctl/visitors/user.py
@@ -3,6 +3,9 @@ from .packages import Packages
 
 class User:
     """Secure the image with a user"""
+    packages = [
+        'shadow',
+    ]
 
     def __init__(self, username, uid, home):
         self.username = username
@@ -10,20 +13,8 @@ class User:
         self.home = home
         self.user_created = False
 
-    def init_build(self, script):
-        """Inject the Packages visitor if necessary."""
-        packages = script.container.visitor('packages')
-        if not packages:
-            index = script.container.visitors.index(self)
-            script.container.visitors.insert(index, Packages())
-
-    def pre_build(self, script):
-        """Inject the shadow package for the usermod command"""
-        if script.container.variable('mgr') == 'apk':
-            script.container.variable('packages').append('shadow')
-
-    def build(self, script):
-        script.append(f'''
+    async def build(self, script):
+        await script.append(f'''
             if {script._run('id ' + str(self.uid))}; then
                 i=$({script._run('id -gn ' + str(self.uid))})
                 {script._run('usermod -d ' + self.home + ' -l ' + self.username + ' $i')}