#!/usr/bin/env python3

'''
Author: Yanrui Hu
Date: 2023/5/22
Description: 使用docker模块进行自动化容器创建与容器内脚本执行、上传文件至容器、commit容器、删除容器、push容器等操作
Usage: 直接运行此脚本，会打印出帮助信息
Keywords: docker, python, automate, container
'''
import docker, os, subprocess, sys, argparse, logging

DESCRIPTION = """通过指定镜像启动容器, 将本地文件上传到容器中, 在容器中执行bash命令,
    并将容器commit得到镜像, 最后推送到本地镜像仓库。
        注意：有几点需要说明：
        1. 如果没有安装docker, 会自动下载并安装docker
        2. 如果执行此脚本的用户没有root权限, 会自动创建docker用户组, 并将当前用户加入docker用户组
        3. 如果没有启动本地镜像仓库, 会自动启动本地镜像仓库 (通过提供-l来指定)
        4. -i IMAGE 对应于 docker run 的参数
        5. -o OUT_IMAGE, -a AUTHOR, -m MESSAGE 对应于 docker commit 的参数
        6. -s SRC, -d DEST 对应于 docker cp 的参数
        8. -c CMD, -u USER, -w WORKDIR 对应于 docker exec 的参数
        """
EPILOG = """Example:
    python3 automate.py -i ubuntu:20.04 -o localhost:5000/ubuntu:v1.0 \
        -c cmd_file.txt\
        -a 'HYR <yanruinku@qq.com>' -m "My ubuntu image, version 1.0"\
        -l
    # Tips: 打开此文件，阅读头部的注释，可以了解更多细节
    """

# ------------------------------------------------------------
# 配置日志
# ------------------------------------------------------------
# 创建一个logger
running_logger = logging.getLogger("running")
running_logger.setLevel(logging.DEBUG)

# 创建一个handler，用于写入日志文件
file_handler = logging.FileHandler("running.log")
file_handler.setLevel(logging.DEBUG)

# 设置日志格式
formatter = logging.Formatter(
    '%(asctime)s - %(name)s - %(levelname)s - %(filename)s:%(lineno)d %(funcName)s - %(message)s'
)

file_handler.setFormatter(formatter)

# 给logger添加handler
running_logger.addHandler(file_handler)


def get_all_containers_status():
    """Get all containers' status"""
    client = docker.from_env()
    for container in client.containers.list(all=True):
        print(f'container: {container.short_id} Status: {container.status}')


def get_client_events():
    """Get real-time events from the server. Similar to the `docker events` command."""
    client = docker.from_env()
    for event in client.events(decode=True):
        print(event)


def run_and_exec(image: str, user: str, workdir: str, cmd_file: str) -> str:
    """通过给定的镜像启动一个容器，并执行 cmd_file 中的命令, 返回容器的短ID
    Params:
        - image: 镜像名
        - user: 容器内执行命令的用户
        - workdir: 容器内执行命令的工作目录
        - cmd_file: 存放命令的文件，可以是 .txt 或 .sh 文件
    Return: 容器的短ID
    """
    running_logger.info(f"run_and_exec: {locals()}")

    client = docker.from_env()
    container = client.containers.run(
        image=image,
        command='/bin/bash',  # 启动一个bash shell 可以一直等待下去，适合让容器保持后台执行不退出
        detach=True,
        tty=True,  # 分配一个伪终端，一定要设置为True，否则，容器进程会自动结束，从而容器被remove
        stdin_open=False,  # 不要打开标准输入! 否则，容器进程从stdin读不到数据会自动结束，从而容器被remove
        remove=True,
        network='host',  # 用于配置网络
    )  # 启动一个容器在后台运行，备用

    print(f'Container {container.short_id} started.')
    running_logger.info(f'Container {container.short_id} started.')

    # 判断 cmd_file 是否存在
    if not os.path.exists(cmd_file):
        print(f"Error6: {cmd_file} not found.")
        exit(1)

    running_logger.info(f"cmd_file: {cmd_file}")
    # 判断 cmd_file 的后缀名是 txt 还是 sh
    if cmd_file.endswith(".txt"):
        running_logger.info(f"==== Starting shell commands execution ====")

        with open(cmd_file, "r") as f:
            for cmd in f.readlines():
                cmd = cmd.strip()
                # 如果读到文件末尾，则退出

                if not cmd or cmd.startswith("#"):
                    continue

                running_logger.info(f"cmd: {cmd}")

                print(f"{user}@:{container.short_id}# cmd")
                running_logger.info(f"{user}@:{container.short_id}# {cmd}")

                command = ["/bin/bash", "-xc", cmd] # TODO: 或许将来会修改这里的参数
                try:
                    """ 这是老代码，脚本执行过程中的信息不能及时看到，不好用
                    res = container.exec_run(command, user=user, workdir=workdir)
                    running_logger.info(f"res.output: {res.output.decode()}")
                    print(f"命令在容器内执行结果的输出:\n{res.output.decode()}")

                    running_logger.info(f"res.exit_code: {res.exit_code}")
                    print(f"命令在容器内执行的exit code: {res.exit_code}") """

                    res = subprocess.run(f'docker exec -it {container.short_id}'+ " ".join(command), shell=True)
                    exit_code = res.returncode
                    running_logger.info(f"Exitcode of {cmd} is: {exit_code}")

                except docker.errors.APIError as e:
                    running_logger.info(f"Exception: {e}")
                    print(f"Exception: {e}")
                    exit(2)

                # 如果commmand执行的返回值不为0，则打印错误信息并退出
                assert exit_code == 0, f"Error2: exit_code不为0 请检查命令是否正确"

        running_logger.info(f"==== Ending shell commands execution ====")
    elif cmd_file.endswith(".sh"):
        # 将 cmd_file.sh 打包成 cmd_file.tar之后 上传到容器内
        cmd = f"tar -cf cmd_file.tar '{cmd_file}'"
        assert subprocess.run(cmd, shell=True).returncode == 0
        running_logger.info(f"Succeed: 打包文件 {cmd_file} 成 cmd_file.tar 成功")

        with open("cmd_file.tar", "rb") as f:
            assert container.put_archive("/root", f.read()), "Error5: 上传文件cmd_file.tar失败"
        running_logger.info(f"Succeed: 上传文件cmd_file.tar成功")

        # 清理本地的 cmd_file.tar
        cmd = "rm -rf cmd_file.tar"
        assert subprocess.run(cmd, shell=True).returncode == 0
        running_logger.info(f"Succeed: 清理本地的 cmd_file.tar 成功")

        running_logger.info(f"==== Starting shell script execution ====")

        command = ["/bin/bash", '-x', f"'/root/{cmd_file}'"] # TODO: 或许将来会修改这里的bash参数
        try:
            """ 这是老代码，脚本执行过程中的信息不能及时看到，不好用
            res = container.exec_run(command, user=user, workdir=workdir)
            running_logger.info(f"res.output: {res.output.decode()}")
            print(f"脚本在容器内执行结果的输出:\n{res.output.decode()}")

            running_logger.info(f"res.exit_code: {res.exit_code}")
            print(f"脚本在容器内执行的exit code: {res.exit_code}") """

            # 新代码
            res = subprocess.run(f"docker exec -it {container.short_id} /bin/bash -x '/root/{cmd_file}'", shell=True)
            exit_code = res.returncode
            running_logger.info(f"Exitcode of {cmd_file} is: {exit_code}")


        except docker.errors.APIError as e:
            running_logger.info(f"Exception: {e}")
            print(f"Exception: {e}")
            exit(3)

        # 如果commmand执行的返回值不为0，则打印错误信息并退出
        assert exit_code == 0, f"Error3: exit_code不为0 请检查脚本是否有错误"
        running_logger.info(f"==== Ending shell commands execution ====")
    else:
        running_logger.info(f"Error4: {cmd_file} is not a .txt or .sh file.")
        print(f"Error4: {cmd_file} is not a .txt or .sh file.")
        exit(4)

    running_logger.info(f"Succeed: {container.short_id} 执行 {cmd_file}成功")
    print(f"Command_file {cmd_file} executed completely.")
    return container.short_id


def put_file_to_container(container_id: str, src: str, dest: str):
    """将本地文件 src 上传到容器 container_id 中的 dst 目录下
    Params:
        - container_id: 容器的短ID
        - src: 本地文件的路径
        - dst: 容器内的目标路径
    """
    running_logger.info(f"put_file_to_container: {locals()}")

    client = docker.from_env()
    container = client.containers.get(container_id)

    # 判断 src 是否存在
    if not os.path.exists(src):
        running_logger.info(f"Error7: {src} not found.")
        print(f"Error7: {src} not found.")
        exit(4)

    # 必须先将 src 打包成 tar 文件，再上传到容器中
    assert subprocess.run(f'tar -cf src.tar {src}', shell=True).returncode == 0

    with open("src.tar", "rb") as f:
        assert container.put_archive(dest, f.read()), "Error8: 上传文件src.tar失败"

    # 清理本地的 src.tar
    assert subprocess.run("rm src.tar", shell=True).returncode == 0

    running_logger.info(f"Succeed: 上传文件 {src} 到容器 {container_id} 的 {dest} 成功")
    print(f"File {src} uploaded to {dest}.")


def commit_and_push(container_id: str, out_image: str, author: str, message: str):
    """将容器commit得到镜像out_image, 并推送到镜像仓库
    Params:
        - container_id: 容器的短ID
        - out_image: 输出的镜像名
        - author: 镜像作者
        - message: 镜像的描述信息
    """
    running_logger.info(f"commit_and_push: {locals()}")

    client = docker.from_env()
    container = client.containers.get(container_id)

    repo, tag = out_image.rsplit(":", 1)

    # 将容器commit得到镜像
    # container.commit(repo, tag, author=author, message=message)
    res = subprocess.run(f"docker commit -a='{author}' -m='{message}' {container_id} '{out_image}'", shell=True)
    assert res.returncode == 0, f"Error9: 镜像 {out_image} commit失败"
    running_logger.info(f"Succeed: 镜像 {out_image} commit成功")
    print(f"Image {out_image} commited.")

    # 停掉容器
    container.stop()  # run 时，指定了 remove=True，所以这里不需要手动删除容器

    # 将镜像推送到镜像仓库
    # res = client.api.push(repo, tag, stream=False, decode=True)
    # running_logger.info(f"res: {res}")
    res = subprocess.run(f"docker push {out_image}", shell=True)
    assert res.returncode == 0, f"Error10: 镜像 {out_image} push失败"
    print(f"Image {out_image} pushed.")
    running_logger.info(f"Succeed: 镜像 {out_image} push成功")


def ensure_docker():
    """确保 docker的存在, 下载并安装 docker"""
    cmd1 = "curl -fsSL get.docker.com -o get-docker.sh"
    cmd2 = "sh get-docker.sh --mirror Aliyun"
    if os.geteuid() != 0:
        cmd2 = "sudo " + cmd2

    print("正在检测 docker 是否已安装……")

    res = subprocess.run("which docker", shell=True, capture_output=True)
    if res.returncode != 0:
        print("docker 尚未安装，正在安装……")
        res = subprocess.run(cmd1, shell=True, capture_output=True)
        assert res.returncode == 0, print(f"{cmd1}失败！ 请自行下载安装脚本")

        res = subprocess.run(cmd2, shell=True, capture_output=True)
        assert res.returncode == 0, print(f"{cmd2}失败! 请手动排查错误并安装")
        print("docker 安装成功 ✅")
    else:
        print("docker 已安装 ✅")

    # 检查docker后台服务是否可以使用
    cmd = "docker info > /dev/null 2>&1"
    res = subprocess.run(cmd, shell=True, capture_output=True)
    if res.returncode == 0:
        return

    # 启动 docker
    cmd3 = "systemctl enable docker"
    cmd4 = "systemctl start docker"

    if os.geteuid() != 0:
        cmd3 = "sudo " + cmd3
        cmd4 = "sudo " + cmd4

    print("正在启动 Docker 后台服务")
    res = subprocess.run(cmd3, shell=True, capture_output=True)
    assert res.returncode == 0, print(f"{cmd3}失败! 请手动排查错误")

    res = subprocess.run(cmd4, shell=True, capture_output=True)
    assert res.returncode == 0, print(f"{cmd4}失败! 请手动排查错误")


def create_docker_grp():
    """为非 root 用户创建 docker 用户组, 并加入其中"""
    if os.geteuid() == 0:
        running_logger.info("当前用户为 root 用户，无需创建 docker 用户组")
        return  # root 用户不必创建 docker 用户组
    cmd1 = "sudo groupadd docker"
    cmd2 = "sudo usermod -aG docker $USER"

    # res = os.popen("cat /etc/group | grep docker").read()
    cmd = "cat /etc/group | grep docker"
    res = subprocess.run(cmd, shell=True, capture_output=True)
    assert res.returncode == 0, print(f"{cmd}失败! 请手动排查错误")

    # 检测 /etc/group 文件 是否有 docker 组
    if "docker" in res.stdout.decode():
        print("docker 用户组已存在 ✅")
    else:
        res1 = subprocess.run(cmd1, shell=True, capture_output=True)
        assert res1.returncode == 0, print(f"{cmd1}失败! 请手动排查错误")
        del cmd1, res1
    # 获取当前用户名
    username = os.getlogin()

    # 检测当前用户是否在 docker 组中
    if username in res.stdout.decode():
        print("当前登录用户已在 docker 用户组 ✅")
    else:
        res2 = subprocess.run(cmd2, shell=True, capture_output=True)
        assert res2.returncode == 0, print(f"{cmd2}失败! 请手动排查错误")
        del res2


def ensure_local_registry():
    """确保本地镜像存储库已启动"""
    running_logger.info("ensure_local_registry")
    cmd = "docker ps -a | grep registry"
    res = subprocess.run(cmd, shell=True, capture_output=True)
    if res.returncode == 0 and "Up" in res.stdout.decode():
        return
    elif res.returncode == 0 and "Exited" in res.stdout.decode():
        cmd = "docker start registry"
        res = subprocess.run(cmd, shell=True, capture_output=True)
        assert res.returncode == 0, print(f"{cmd}失败! 请手动排查错误")
        return

    # 启动本地镜像存储库
    cmd = "docker run -d -p 5000:5000 --restart=always --name registry registry"
    res = subprocess.run(cmd, shell=True, capture_output=True)
    assert res.returncode == 0, print(f"{cmd}失败! 请手动排查错误")


def main():
    """程序启动入口，解析命令行参数，执行相应操作"""
    parser = argparse.ArgumentParser(description=DESCRIPTION, epilog=EPILOG)
    parser.add_argument(
        "-i",
        "--image",
        type=str,
        required=True,
        help="用于启动容器的镜像名，格式为 <仓库名>:<标签>，如 docker.io/ubuntu:latest",
    )
    parser.add_argument(
        "-o",
        "--out-image",
        type=str,
        required=True,
        help="得到的镜像名，格式为 <仓库名>:<标签>，如 localhost:5000/ubuntu:v1.0",
    )
    parser.add_argument(
        "-s",
        "--src",
        type=str,
        required=False,
        help="欲上传的本地文件路径，如 /home/xxx/xxx, 支持文件或目录, 必须与 -d 参数同时使用",
    )
    parser.add_argument(
        "-d",
        "--dest",
        type=str,
        required=False,
        default="/root",
        help="上传至容器内部的路径，如 /root/xxx/xxx, 必须与 -s 参数同时使用",
    )
    parser.add_argument(
        "-c",
        "--command-file",
        type=str,
        required=True,
        help="欲在容器中执行的命令文件，可以是.txt或.sh文件",
    )
    parser.add_argument(
        "-u",
        "--user",
        type=str,
        required=False,
        default="root",
        help="容器中执行命令的用户，如 root",
    )
    parser.add_argument(
        "-w",
        "--workdir",
        type=str,
        required=False,
        default="/root",
        help="容器中执行命令的工作目录，如 /root",
    )
    parser.add_argument(
        "-a",
        "--author",
        type=str,
        required=False,
        help="镜像作者，如 xxx",
    )
    parser.add_argument(
        "-m",
        "--message",
        type=str,
        required=False,
        help="镜像描述信息，如 xxx",
    )
    parser.add_argument(
        "-l",
        "--local-registry",
        action="store_true",
        required=False,
        help="启动一个本地镜像存储库",
    )

    if len(sys.argv) == 1:  # 如果没有提供任何参数, 则打印帮助信息
        parser.print_help()
        sys.exit()

    args = parser.parse_args()
    running_logger.info(args)
    print(f'{vars(args)=}')

    # 确保 docker 的存在
    ensure_docker()

    # 为非 root 用户创建 docker 用户组, 并加入其中
    create_docker_grp()

    # 启动本地镜像存储库,若需要
    if args.local_registry:
        ensure_local_registry()

    container_id = run_and_exec(
        image=args.image,
        user=args.user,
        workdir=args.workdir,
        cmd_file=args.command_file,
    )
    assert container_id, print("获取container_id失败! 请手动排查错误")
    if args.src:
        put_file_to_container(container_id=container_id, src=args.src, dst=args.dest)

    # 将容器commit得到镜像，然后推送到镜像仓库
    commit_and_push(container_id, args.out_image, args.author, args.message)


if __name__ == "__main__":
    main()