Merge "servo: For all commands that read stdout, decode" into main
diff --git a/scripts/custom_os_image_analysis.py b/scripts/custom_os_image_analysis.py
new file mode 100644
index 0000000..5849b80
--- /dev/null
+++ b/scripts/custom_os_image_analysis.py
@@ -0,0 +1,83 @@
+import os
+from os.path import expanduser
+import sys
+
+from google.api_core import page_iterator
+from google.cloud import storage
+
+"""
+This is a standalone script used to get the number of custom builds 
+that have been uploaded to the moblab buckets. This analysis helps with 
+figuring out who are the active users of the Custom OS Image feature. 
+This script is expected to be run outside the Moblab dev environment
+in the local setup. 
+This script is expected to be run by members of moblab team who
+have full access to the GCP account. 
+Download the account login json from the GCP and put it in ~/.gsutil/moblab_login
+ 
+"""
+# download the service_account json from the GCP web portal
+# Set environment variable
+path_where_key_is_stored = expanduser(
+    "~/.gsutil/moblab_login/chromeos-partner-moblab-account-login.json"
+)
+os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = path_where_key_is_stored
+
+
+def _item_to_value(iterator, item):
+    return item
+
+
+def list_directories(gcs, bucket_name, prefix):
+    if prefix and not prefix.endswith("/"):
+        prefix += "/"
+    extra_params = {"projection": "noAcl", "prefix": prefix, "delimiter": "/"}
+    path = "/b/" + bucket_name + "/o"
+    iterator = page_iterator.HTTPIterator(
+        client=gcs,
+        api_request=gcs._connection.api_request,
+        path=path,
+        items_key="prefixes",
+        item_to_value=_item_to_value,
+        extra_params=extra_params,
+    )
+    return [x for x in iterator]
+
+
+def main(argv):
+    client = storage.Client()
+    # get the list of buckets under the account
+    buckets = client.list_buckets()
+    bucket_names = [x.name for x in buckets]
+    bucket_level_info = {}
+
+    # for each bucket, figure out the number of custom_builds
+    for bucket_name in bucket_names:
+        folders = list_directories(client, bucket_name, "")
+        num_builds_per_partner = 0
+        for folder in folders:
+            if "-local" in folder:
+                # figure out the list of custom OS builds uploaded in this folder
+                blobs = client.list_blobs(
+                    bucket_name, prefix=folder, delimiter=""
+                )
+                # always expect something after the first "/" because of the above query
+                build_folders = [blob.name.split("/")[1] for blob in blobs]
+                # Example of build_folder:
+                # chromeos-moblab-peng-staging/eve-local/R102-14679.0.0-d2022_04_11_1033
+                # make build_folder names unique
+                uniq_builds = set(build_folders)
+                # remove dummy results in the result
+                if "" in uniq_builds:
+                    uniq_builds.remove("")
+                num_builds_per_partner += len(uniq_builds)
+        if num_builds_per_partner:
+            print(
+                f"{bucket_name} bucket has {num_builds_per_partner} custom builds"
+            )
+            bucket_level_info[bucket_name] = num_builds_per_partner
+    print(bucket_level_info)
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])