From 90c95e3a8ddc9f119dfb51008536c14c5dfe7257 Mon Sep 17 00:00:00 2001
From: Farid Zakaria <fmzakari@google.com>
Date: Sat, 6 Apr 2024 04:36:13 +0000
Subject: [PATCH] Change docker2sqlelf to be in Python

Rewrite docker2sqlelf to be written in Python rather than bash for
reusability.
---
 tools/docker2sqlelf | 166 +++++++++++++++++++++++---------------------
 1 file changed, 88 insertions(+), 78 deletions(-)

diff --git a/tools/docker2sqlelf b/tools/docker2sqlelf
index 376c196..32bae06 100755
--- a/tools/docker2sqlelf
+++ b/tools/docker2sqlelf
@@ -1,78 +1,88 @@
-#!/usr/bin/env bash
-
-set -e -o pipefail
-
-# Default behavior is to clean up the temporary directory
-KEEP_TEMP_DIR=false
-
-# Parse command-line arguments
-while [[ "$#" -gt 0 ]]; do
-    case $1 in
-        -k|--keep) KEEP_TEMP_DIR=true ;;
-        *) IMAGE_NAME="$1" ;;
-    esac
-    shift
-done
-
-# Ensure an image name was provided
-if [ -z "$IMAGE_NAME" ]; then
-    echo "Usage: $0 [-k|--keep] <docker-image-name>"
-    exit 1
-fi
-
-# Create a temporary directory for the container's filesystem
-TEMP_DIR=$(mktemp -d)
-echo "Created temporary directory at $TEMP_DIR"
-
-# Function to cleanup on exit, if not keeping the temporary directory
-cleanup() {
-    if [ "$KEEP_TEMP_DIR" = false ]; then
-        echo "Cleaning up..."
-        # Remove temporary directory
-        rm -rf "$TEMP_DIR"
-        echo "Removed temporary directory $TEMP_DIR"
-    else
-        echo "Keeping temporary directory $TEMP_DIR"
-    fi
-}
-
-# Create a temporary directory for the container's filesystem
-TEMP_DIR=$(mktemp -d)
-echo "Created temporary directory at $TEMP_DIR"s
-
-# Trap EXIT to cleanup
-trap cleanup EXIT
-
-# Create docker container from the image
-CONTAINER_ID=$(docker create "$IMAGE_NAME")
-echo "Created container with ID $CONTAINER_ID"
-
-# Export the container's filesystem and untar it to the temporary directory
-docker export "$CONTAINER_ID" | tar -C "$TEMP_DIR" -xf -
-echo "Exported and extracted container's filesystem to $TEMP_DIR"
-
-# Remove the docker container
-docker rm "$CONTAINER_ID"
-echo "Removed container $CONTAINER_ID"
-
-# Replace colons with hyphens
-MODIFIED_IMAGE_NAME="${IMAGE_NAME//:/-}"
-
-# Uncomment the below lines if this script is failing on a distribution
-# it will run sqlelf on each file one at a time, and it will tell you
-# which file it is failing on
-# echo "Running sqlelf on the contents of $TEMP_DIR one file at a time."
-# KEEP_TEMP_DIR=true
-# find "$TEMP_DIR" -type f -print0 | while IFS= read -r -d $'\0' file; do
-#    if ! file "$file" | grep -q 'ELF'; then
-#        continue
-#    fi
-#    echo "Processing $file with sqlelf"
-#    sqlelf "$file" --sql "SELECT 1;"
-# done
-
-# Run sqlelf tool on the contents of the temporary directory
-echo "Running sqlelf on the contents of $TEMP_DIR"
-sqlelf "$TEMP_DIR" --sql ".backup ${MODIFIED_IMAGE_NAME}.sqlite"
-
-echo "Script completed successfully."
+#!/usr/bin/env python3
+
+import argparse
+import atexit
+import os
+import shutil
+import tempfile
+from functools import reduce
+
+import docker
+
+from sqlelf import elf, sql
+
+
+def docker2sqelf(image_name: str, keep_temp_dir: bool = False) -> str:
+    client = docker.from_env()
+
+    temp_dir = tempfile.mkdtemp()
+    print(f"Created temporary directory at {temp_dir}")
+
+    def cleanup():
+        if not keep_temp_dir:
+            print("Cleaning up...")
+            shutil.rmtree(temp_dir)
+            print(f"Removed temporary directory {temp_dir}")
+        else:
+            print(f"Keeping temporary directory {temp_dir}")
+
+    atexit.register(cleanup)
+
+    container = client.containers.create(image_name)
+    print(f"Created container with ID {container.id}")
+
+    export_path = f"{temp_dir}/container.tar"
+    with open(export_path, "wb") as out_f:
+        bits = container.export()
+        for chunk in bits:
+            out_f.write(chunk)
+    print(f"Exported container's filesystem to {export_path}")
+
+    shutil.unpack_archive(export_path, temp_dir)
+    print(f"Extracted container's filesystem to {temp_dir}")
+
+    container.remove()
+    print(f"Removed container {container.id}")
+
+    modified_image_name = image_name.replace(":", "-")
+
+    filenames: list[str] = reduce(
+        lambda a, b: a + b,
+        map(
+            lambda dir: (
+                [
+                    os.path.join(root, file)
+                    for root, _, files in os.walk(dir)
+                    for file in files
+                ]
+                if os.path.isdir(dir)
+                else [dir]
+            ),
+            [temp_dir],
+        ),
+    )
+
+    filenames = [f for f in filenames if os.path.isfile(f)]
+
+    print("Creating sqlelf database")
+    engine = sql.make_sql_engine(filenames, elf.CacheFlag.ALL())
+
+    print("Dumping the sqlite database")
+    database_filename = f"{modified_image_name}.sqlite"
+    engine.dump(database_filename)
+
+    print(f"Created database {database_filename}")
+    return database_filename
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Convert docker image to sqlelf database."
+    )
+    parser.add_argument("image_name", help="Docker image name")
+    parser.add_argument(
+        "-k", "--keep", help="Keep temporary directory", action="store_true"
+    )
+    args = parser.parse_args()
+
+    docker2sqelf(args.image_name, args.keep)