Automatically extract downloaded blobs

1 files changed, 28 insertions, 7 deletions

diff --git a/src/main/java/io/github/jshipit/BlobDownloader.java b/src/main/java/io/github/jshipit/BlobDownloader.java
index 6a17057..91f126e 100644
--- a/src/main/java/io/github/jshipit/BlobDownloader.java
+++ b/src/main/java/io/github/jshipit/BlobDownloader.java
@@ -1,11 +1,10 @@
 package io.github.jshipit;
 
 import me.tongfei.progressbar.ProgressBar;
+import org.apache.commons.compress.archivers.tar.*;
+import org.apache.commons.compress.compressors.gzip.GzipCompressorInputStream;
 
-import java.io.FileOutputStream;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
+import java.io.*;
 import java.net.*;
 
 public class BlobDownloader extends Thread {
@@ -14,15 +13,17 @@ public class BlobDownloader extends Thread {
     private final String digest;
     private final String[][] headers;
     private final String tmpdir;
+    private final boolean extract;
 
-    public BlobDownloader(String url, String digest, String[][] headers, String tmpdir) {
+    public BlobDownloader(String url, String digest, String[][] headers, String tmpdir, boolean extract) {
         this.url = url;
         this.digest = digest;
         this.headers = headers;
         this.tmpdir = tmpdir;
+        this.extract = extract;
     }
 
-    public void run() throws RuntimeException {
+    public void run() {
         URL url_obj = null;
         try {
             url_obj = new URI(this.url).toURL();
@@ -60,7 +61,7 @@ public class BlobDownloader extends Thread {
         int fileSize = con.getContentLength();
 
         try (InputStream in = con.getInputStream();
-             OutputStream out = new FileOutputStream(tmpdir + "/" + digest.replace("sha256:", "") + "layer.tar")) {
+             OutputStream out = new FileOutputStream(tmpdir + "/" + digest.replace("sha256:", "") + ".tar.gz")) {
             byte[] buffer = new byte[4096];
             int bytesRead;
             try (ProgressBar pb = new ProgressBar("Download blob "+digest.replace("sha256:", ""), fileSize)) {
@@ -73,5 +74,25 @@ public class BlobDownloader extends Thread {
             System.out.println("Failed to download blob "+digest);
             e.printStackTrace();
         }
+
+        File extracted = new File(tmpdir + "/" + digest.replace("sha256:", ""));
+        if (extracted.exists()) {
+            extracted.delete();
+        }
+
+        if (!extract) {
+            return;
+        }
+        TarManager tarManager = new TarManager();
+        try {
+            tarManager.untar(tmpdir + "/" + digest.replace("sha256:", "") + ".tar.gz", extracted);
+        } catch (IOException e) {
+            System.out.println("Failed to extract blob "+digest);
+            e.printStackTrace();
+        }
+        File tar = new File(tmpdir + "/" + digest.replace("sha256:", "") + ".tar.gz");
+        if (tar.exists()) {
+            tar.delete();
+        }
     }
 }