p4runtime-sh-docker

#!/usr/bin/env python3

# Copyright 2019 Barefoot Networks, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import argparse
from collections import namedtuple
import logging
import os.path
import sys
import tempfile
import shutil
import subprocess


def main():
    FwdPipeConfig = namedtuple('FwdPipeConfig', ['p4info', 'bin'])

    def pipe_config(arg):
        try:
            paths = FwdPipeConfig(*[x for x in arg.split(',')])
            if len(paths) != 2:
                raise argparse.ArgumentError
            return paths
        except Exception:
            raise argparse.ArgumentError(
                "Invalid pipeline config, expected <p4info path>,<binary config path>")

    parser = argparse.ArgumentParser(description='P4Runtime shell docker wrapper', add_help=False)
    parser.add_argument('--config',
                        help='If you want the shell to push a pipeline config to the server first',
                        metavar='<p4info path (text)>,<binary config path>',
                        type=pipe_config, action='store', default=None)
    parser.add_argument('-v', '--verbose', help='Increase output verbosity',
                        action='store_true')
    parser.add_argument('--docker-image', help='p4runtime-sh Docker image to use',
                        type=str, action='store', default='p4lang/p4runtime-sh')
    parser.add_argument('--volumes',
                        help='Additonal volumes for the Docker container (same as docker run)',
                        type=str, action='append', metavar='volume', default=[])
    parser.add_argument('--mounts',
                        help='Additonal mounts for the Docker container (same as docker run)',
                        type=str, action='append', metavar='mount', default=[])
    parser.add_argument('--cacert',
                        help='CA certificate to verify peer against, for secure connections',
                        metavar='<path to .pem>',
                        type=str, action='store', default=None)
    parser.add_argument('--cert',
                        help='Path to client certificate, for mutual authentication',
                        metavar='<path to .pem>',
                        type=str, action='store', default=None)
    parser.add_argument('--private-key',
                        help='Path to client private key, for mutual authentication',
                        metavar='<path to .pem>',
                        type=str, action='store', default=None)
    args, unknown_args = parser.parse_known_args()

    docker_args = []
    new_args = []

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
        new_args.append('--verbose')

    files_to_mount = []
    mount_path = "/p4rt_data"

    if args.config is not None:
        fname_p4info = "p4info.pb.txt"
        fname_bin = "config.bin"
        files_to_mount += [(fname_p4info, args.config.p4info)]
        files_to_mount += [(fname_bin, args.config.bin)]
        new_args.extend(["--config", "{},{}".format(
            os.path.join(mount_path, fname_p4info), os.path.join(mount_path, fname_bin))])

    if args.cacert is not None:
        fname_cacert = "cacert.pem"
        files_to_mount += [(fname_cacert, args.cacert)]
        new_args.extend(["--cacert", os.path.join(mount_path, fname_cacert)])

    if args.cert is not None:
        fname_cert = "cert.pem"
        files_to_mount += [(fname_cert, args.cert)]
        new_args.extend(["--cert", os.path.join(mount_path, fname_cert)])

    if args.private_key is not None:
        fname_private_key = "key.pem"
        files_to_mount += [(fname_private_key, args.private_key)]
        new_args.extend(["--private-key", os.path.join(mount_path, fname_private_key)])

    for f in files_to_mount:
        path = f[1]
        if not os.path.isfile(path):
            logging.critical("'{}' is not a valid file".format(path))
            sys.exit(1)

    tmpdir = None
    if len(files_to_mount) > 0:
        logging.debug(
            "Created temporary directory '{}', it will be mounted in the docker as '{}'".format(
                tmpdir, mount_path))
        tmpdir = tempfile.mkdtemp(dir="/tmp")
        docker_args.extend(["-v", "{}:{}".format(tmpdir, mount_path)])

    for f in files_to_mount:
        fname, path = f
        shutil.copy(path, os.path.join(tmpdir, fname))

    for volume in args.volumes:
        docker_args.extend(["-v", volume])
    for mount in args.mounts:
        docker_args.extend(["--mount", mount])

    cmd = ["docker", "run", "-ti", "--rm"]
    cmd.extend(docker_args)
    cmd.append(args.docker_image)
    cmd.extend(new_args)
    cmd.extend(unknown_args)
    logging.debug("Running cmd: {}".format(" ".join(cmd)))

    subprocess.run(cmd)

    if tmpdir is not None:
        logging.debug("Cleaning up...")
        try:
            shutil.rmtree(tmpdir)
        except Exception:
            logging.error("Error when removing temporary directory '{}'".format(tmpdir))


if __name__ == '__main__':
    main()