From a335a9dbbc82341397ae12013c98ad37cf52107b Mon Sep 17 00:00:00 2001
From: Przemyslaw Gorszkowski <przemek.gorszkowski@redembedded.com>
Date: Mon, 26 Apr 2021 10:19:21 +0200
Subject: [PATCH] S3 fetcher: add progress handler for S3 cp command

---
 lib/bb/fetch2/s3.py | 41 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 40 insertions(+), 1 deletion(-)

diff --git a/lib/bb/fetch2/s3.py b/lib/bb/fetch2/s3.py
index ffca73c8..6b8ffd53 100644
--- a/lib/bb/fetch2/s3.py
+++ b/lib/bb/fetch2/s3.py
@@ -18,10 +18,47 @@ The aws tool must be correctly installed and configured prior to use.
 import os
 import bb
 import urllib.request, urllib.parse, urllib.error
+import re
 from bb.fetch2 import FetchMethod
 from bb.fetch2 import FetchError
 from bb.fetch2 import runfetchcmd
 
+def convertToBytes(value, unit):
+    value = float(value)
+    if (unit == "KiB"):
+        value = value*1024.0;
+    elif (unit == "MiB"):
+        value = value*1024.0*1024.0;
+    elif (unit == "GiB"):
+        value = value*1024.0*1024.0*1024.0;
+    return value
+
+class S3ProgressHandler(bb.progress.LineFilterProgressHandler):
+    """
+    Extract progress information from s3 cp output, e.g.:
+    Completed 5.1 KiB/8.8 GiB (12.0 MiB/s) with 1 file(s) remaining
+    """
+    def __init__(self, d):
+        super(S3ProgressHandler, self).__init__(d)
+        # Send an initial progress event so the bar gets shown
+        self._fire_progress(0)
+
+    def writeline(self, line):
+        percs = re.findall(r'^Completed (\d+.{0,1}\d*) (\w+)\/(\d+.{0,1}\d*) (\w+) (\(.+\)) with\s+', line)
+        if percs:
+            completed = (percs[-1][0])
+            completedUnit = (percs[-1][1])
+            total = (percs[-1][2])
+            totalUnit = (percs[-1][3])
+            completed = convertToBytes(completed, completedUnit)
+            total = convertToBytes(total, totalUnit)
+            progress = (completed/total)*100.0
+            rate = percs[-1][4]
+            self.update(progress, rate)
+            return False
+        return True
+
+
 class S3(FetchMethod):
     """Class to fetch urls via 'aws s3'"""
 
@@ -52,7 +89,9 @@ class S3(FetchMethod):
 
         cmd = '%s cp s3://%s%s %s' % (ud.basecmd, ud.host, ud.path, ud.localpath)
         bb.fetch2.check_network_access(d, cmd, ud.url)
-        runfetchcmd(cmd, d)
+
+        progresshandler = S3ProgressHandler(d)
+        runfetchcmd(cmd, d, False, log=progresshandler)
 
         # Additional sanity checks copied from the wget class (although there
         # are no known issues which mean these are required, treat the aws cli
-- 
2.31.1