multithreaded downloads and non docker registries

Runs each layer download on a different thread.
Allows using registries other than registry.docker.io

4 files changed, 289 insertions, 198 deletions

diff --git a/src/main/java/io/github/jshipit/BlobDownloader.java b/src/main/java/io/github/jshipit/BlobDownloader.java
new file mode 100644
index 0000000..5fa3959
--- /dev/null
+++ b/src/main/java/io/github/jshipit/BlobDownloader.java
@@ -0,0 +1,77 @@
+package io.github.jshipit;
+
+import me.tongfei.progressbar.ProgressBar;
+
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.net.*;
+
+public class BlobDownloader extends Thread {
+
+    private final String url;
+    private final String digest;
+    private final String[][] headers;
+    private final String tmpdir;
+
+    public BlobDownloader(String url, String digest, String[][] headers, String tmpdir) {
+        this.url = url;
+        this.digest = digest;
+        this.headers = headers;
+        this.tmpdir = tmpdir;
+    }
+
+    public void run() throws RuntimeException {
+        URL url_obj = null;
+        try {
+            url_obj = new URI(this.url).toURL();
+            System.out.println(url_obj.toString());
+        } catch(URISyntaxException | MalformedURLException e) {
+            System.out.println("URISyntaxException | MalformedURLException");
+            e.printStackTrace();
+        }
+
+        assert url_obj != null;
+        HttpURLConnection con;
+
+        try {
+            con = (HttpURLConnection) url_obj.openConnection();
+            con.setRequestMethod("GET");
+            for (String[] header : this.headers) {
+                con.setRequestProperty(header[0], header[1]);
+            }
+            con.connect();
+
+            if (con.getResponseCode() != 200) {
+                System.out.println("Error: " + con.getResponseCode());
+                throw new RuntimeException("Failed : HTTP error code : "
+                        + con.getResponseCode());
+            } else {
+                System.out.println("Success: " + con.getResponseCode());
+            }
+        } catch (IOException e) {
+            System.out.println("Failed to connect");
+            e.printStackTrace();
+            return;
+        }
+
+
+        int fileSize = con.getContentLength();
+
+        try (InputStream in = con.getInputStream();
+             OutputStream out = new FileOutputStream(tmpdir + "/" + digest.replace("sha256:", "") + "layer.tar")) {
+            byte[] buffer = new byte[4096];
+            int bytesRead;
+            try (ProgressBar pb = new ProgressBar("Download blob "+digest.replace("sha256:", ""), fileSize)) {
+                while ((bytesRead = in.read(buffer)) != -1) {
+                    out.write(buffer, 0, bytesRead);
+                    pb.stepBy(bytesRead);
+                }
+            }
+        } catch (IOException e) {
+            System.out.println("Failed to download blob "+digest);
+            e.printStackTrace();
+        }
+    }
+}
diff --git a/src/main/java/io/github/jshipit/DockerAPIHelper.java b/src/main/java/io/github/jshipit/DockerAPIHelper.java
index dca582e..127f90e 100755
--- a/src/main/java/io/github/jshipit/DockerAPIHelper.java
+++ b/src/main/java/io/github/jshipit/DockerAPIHelper.java
@@ -1,162 +1,158 @@
-package io.github.jshipit;

-

-

-

-import com.fasterxml.jackson.core.JsonProcessingException;

-import com.fasterxml.jackson.databind.*;

-import com.fasterxml.jackson.databind.node.ArrayNode;

-

-import java.io.*;

-import java.net.*;

-import java.util.Scanner;

-

-public class DockerAPIHelper {

-

-    private String apiToken;

-    private String apiRepo;

-    private String authURL;

-    private String repository;

-    private String image;

-

-    private String tag;

-

-    public DockerAPIHelper(String apiRepo, String authURL, String repository, String image, String tag) {

-        System.out.println("DockerAPIHelper constructor");

-        this.apiRepo = apiRepo;

-        this.repository = repository;

-        this.authURL = authURL;

-        this.image = image;

-        this.tag = tag;

-        try {

-            apiToken = generateAPIToken();

-        } catch (IOException | RuntimeException e) {

-            System.out.println("IOException | RuntimeException");

-            e.printStackTrace();

-        }

-    }

-

-    public String generateAPIToken() throws IOException, RuntimeException {

-        URL url_obj = null;

-        try {

-            url_obj = new URI(this.authURL + "?scope=repository:" + this.repository + "/" + this.image + ":pull"  + "&service=" + this.apiRepo).toURL();

-            System.out.println(url_obj.toString());

-        } catch (URISyntaxException | MalformedURLException e) {

-            System.out.println("URISyntaxException | MalformedURLException");

-            e.printStackTrace();

-        }

-

-

-        assert url_obj != null;

-        HttpURLConnection con = (HttpURLConnection) url_obj.openConnection();

-        con.setRequestMethod("GET");

-        con.connect();

-        if (con.getResponseCode() != 200) {

-            System.out.println("Error: " + con.getResponseCode());

-            throw new RuntimeException("Failed : HTTP error code : "

-                    + con.getResponseCode());

-        } else {

-            System.out.println("Success: " + con.getResponseCode());

-            String output = "";

-            Scanner scanner = new Scanner(url_obj.openStream());

-            while (scanner.hasNext()) {

-                output += scanner.nextLine();

-            }

-            scanner.close();

-            System.out.println(output);

-            ObjectMapper mapper = new ObjectMapper();

-            JsonNode token = mapper.readTree(output);

-            System.out.println(token.get("token").asText());

-            return token.get("token").asText();

-        }

-    }

-

-    public JsonNode fetchManifestJson() throws IOException, RuntimeException {

-        URL url_obj = null;

-        try {

-            String repo = this.apiRepo.replace("registry", "registry-1");

-            url_obj = new URI("https://" + repo + "/v2/" + this.repository + "/" + this.image + "/manifests/"+this.tag).toURL();

-            System.out.println(url_obj.toString());

-        } catch(URISyntaxException | MalformedURLException e) {

-            System.out.println("URISyntaxException | MalformedURLException");

-            e.printStackTrace();

-        }

-

-        assert url_obj != null;

-        HttpURLConnection con = (HttpURLConnection) url_obj.openConnection();

-        con.setRequestMethod("GET");

-        con.setRequestProperty("Authorization", "Bearer " + this.apiToken);

-        con.setRequestProperty("Accept", "application/vnd.docker.distribution.manifest.v2+json");

-        con.connect();

-

-        if (con.getResponseCode() != 200) {

-            System.out.println("Error: " + con.getResponseCode());

-            throw new RuntimeException("Failed : HTTP error code : "

-                    + con.getResponseCode());

-        } else {

-            System.out.println("Success: " + con.getResponseCode());

-            BufferedReader in = new BufferedReader(new InputStreamReader(con.getInputStream()));

-            String inputLine;

-            StringBuffer content = new StringBuffer();

-            while ((inputLine = in.readLine()) != null) {

-                content.append(inputLine);

-            }

-            in.close();

-            con.disconnect();

-            ObjectMapper mapper = new ObjectMapper();

-            JsonNode manifest = mapper.readTree(content.toString());

-            return manifest;

-        }

-    }

-

-    public void fetchBlob(String digest, String tmpdir) throws IOException, RuntimeException {

-        URL url_obj = null;

-        try {

-            String repo = this.apiRepo.replace("registry", "registry-1");

-            url_obj = new URI("https://" + repo + "/v2/" + this.repository + "/" + this.image + "/blobs/"+digest).toURL();

-            System.out.println(url_obj.toString());

-        } catch(URISyntaxException | MalformedURLException e) {

-            System.out.println("URISyntaxException | MalformedURLException");

-            e.printStackTrace();

-        }

-

-        assert url_obj != null;

-        HttpURLConnection con = (HttpURLConnection) url_obj.openConnection();

-        con.setRequestMethod("GET");

-        con.setRequestProperty("Authorization", "Bearer " + this.apiToken);

-        con.setRequestProperty("Accept", "application/vnd.docker.distribution.manifest.v2+json");

-        con.connect();

-

-        if (con.getResponseCode() != 200) {

-            System.out.println("Error: " + con.getResponseCode());

-            throw new RuntimeException("Failed : HTTP error code : "

-                    + con.getResponseCode());

-        } else {

-            System.out.println("Success: " + con.getResponseCode());

-        }

-        // Stream file download and output progress

-        InputStream inputStream = con.getInputStream();

-        FileOutputStream outputStream = new FileOutputStream(tmpdir+"/"+digest.replace("sha256:", "")+"layer.tar");

-

-        byte[] buffer = new byte[1024];

-        int bytesRead;

-        while ((bytesRead = inputStream.read(buffer)) != -1) {

-            outputStream.write(buffer, 0, bytesRead);

-        }

-

-        outputStream.close();

-        inputStream.close();

-        System.out.println("File downloaded successfully.");

-    }

-

-    public String getApiToken() {

-        return apiToken;

-    }

-

-    public String getImage() {

-        return image;

-    }

-

-    public String getTag() {

-        return tag;

-    }

-}

+package io.github.jshipit;
+
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import me.tongfei.progressbar.*;
+
+import java.io.*;
+import java.net.*;
+import java.sql.Blob;
+import java.util.List;
+import java.util.Map;
+import java.util.Scanner;
+
+public class DockerAPIHelper {
+
+    private String apiToken;
+    private final String apiRepo;
+    private String authURL;
+    private String authService;
+    private final String repository;
+    private final String image;
+    private final String tag;
+
+    public DockerAPIHelper(String apiRepo, String repository, String image, String tag) {
+        System.out.println("DockerAPIHelper constructor");
+        if (apiRepo.contains("registry.docker.io")) {
+            apiRepo = apiRepo.replace("registry", "registry-1"); // Docker is just funny like that
+        }
+        this.apiRepo = apiRepo;
+        this.repository = repository;
+        this.image = image;
+        this.tag = tag;
+        this.authURL = "https://auth.docker.io/token";
+        this.authService = "registry.docker.io";
+        try {
+            getAuthenticationUrl();
+            apiToken = generateAPIToken();
+        } catch (IOException | RuntimeException e) {
+            System.out.println("IOException | RuntimeException");
+            e.printStackTrace();
+        }
+    }
+
+    public void getAuthenticationUrl() throws IOException {
+        URL url_obj = null;
+        try {
+            url_obj = new URI("https://"+this.apiRepo+"/v2/").toURL();
+            System.out.println(url_obj.toString());
+        } catch (URISyntaxException | MalformedURLException e) {
+            e.printStackTrace();
+        }
+
+        assert url_obj != null;
+        HttpURLConnection con = (HttpURLConnection) url_obj.openConnection();
+        con.setRequestMethod("GET");
+        con.connect();
+        if (con.getResponseCode() == 401 ) {
+            Map<String, List<String>> headers = con.getHeaderFields();
+            List<String> authenticate = headers.get("Www-Authenticate");
+            this.authURL = authenticate.get(0).replace("Bearer realm=", "").replace("\"", "").split(",")[0];
+            this.authService = authenticate.get(0).replace("service=", "").replace("\"", "").split(",")[1];
+        } else {
+            this.authURL = "https://auth.docker.io/token";
+            this.authService = "registry.docker.io";
+        }
+    }
+
+    public String generateAPIToken() throws IOException, RuntimeException {
+        URL url_obj = null;
+        try {
+            url_obj = new URI(this.authURL + "?scope=repository:" + this.repository + "/" + this.image + ":pull"  + "&service=" + this.authService).toURL();
+            System.out.println(url_obj.toString());
+        } catch (URISyntaxException | MalformedURLException e) {
+            System.out.println("URISyntaxException | MalformedURLException");
+            e.printStackTrace();
+        }
+
+
+        assert url_obj != null;
+        HttpURLConnection con = (HttpURLConnection) url_obj.openConnection();
+        con.setRequestMethod("GET");
+        con.connect();
+        if (con.getResponseCode() != 200) {
+            System.out.println("Error: " + con.getResponseCode());
+            throw new RuntimeException("Failed : HTTP error code : "
+                    + con.getResponseCode());
+        } else {
+            System.out.println("Success: " + con.getResponseCode());
+            StringBuilder output = new StringBuilder();
+            Scanner scanner = new Scanner(url_obj.openStream());
+            while (scanner.hasNext()) {
+                output.append(scanner.nextLine());
+            }
+            scanner.close();
+            System.out.println(output);
+            ObjectMapper mapper = new ObjectMapper();
+            JsonNode token = mapper.readTree(output.toString());
+            System.out.println(token.get("token").asText());
+            return token.get("token").asText();
+        }
+    }
+
+    public JsonNode fetchManifestJson() throws IOException, RuntimeException {
+        URL url_obj = null;
+        try {
+            url_obj = new URI("https://" + this.apiRepo + "/v2/" + this.repository + "/" + this.image + "/manifests/"+this.tag).toURL();
+            System.out.println(url_obj.toString());
+        } catch(URISyntaxException | MalformedURLException e) {
+            System.out.println("URISyntaxException | MalformedURLException");
+            e.printStackTrace();
+        }
+
+        assert url_obj != null;
+        HttpURLConnection con = (HttpURLConnection) url_obj.openConnection();
+        con.setRequestMethod("GET");
+        con.setRequestProperty("Authorization", "Bearer " + this.apiToken);
+        con.setRequestProperty("Accept", "application/vnd.docker.distribution.manifest.v2+json");
+        con.connect();
+
+        if (con.getResponseCode() != 200) {
+            System.out.println("Error: " + con.getResponseCode());
+            throw new RuntimeException("Failed : HTTP error code : "
+                    + con.getResponseCode());
+        } else {
+            System.out.println("Success: " + con.getResponseCode());
+            BufferedReader in = new BufferedReader(new InputStreamReader(con.getInputStream()));
+            String inputLine;
+            StringBuilder content = new StringBuilder();
+            while ((inputLine = in.readLine()) != null) {
+                content.append(inputLine);
+            }
+            in.close();
+            con.disconnect();
+            ObjectMapper mapper = new ObjectMapper();
+            return mapper.readTree(content.toString());
+        }
+    }
+
+    public void fetchBlob(String digest, String tmpdir) throws IOException, RuntimeException {
+        String url = "https://" + this.apiRepo + "/v2/" + this.repository + "/" + this.image + "/blobs/"+digest;
+        String[][] headers = {{"Authorization", "Bearer "+this.apiToken}, {"Accept", "application/vnd/docker.distribution.manifest.v2+json"}};
+        BlobDownloader downloader = new BlobDownloader(url, digest, headers, tmpdir);
+        downloader.start();
+    }
+
+    public String getApiToken() {
+        return apiToken;
+    }
+
+    public String getImage() {
+        return image;
+    }
+
+    public String getTag() {
+        return tag;
+    }
+}
diff --git a/src/main/java/io/github/jshipit/JshipIT.java b/src/main/java/io/github/jshipit/JshipIT.java
new file mode 100755
index 0000000..30065a0
--- /dev/null
+++ b/src/main/java/io/github/jshipit/JshipIT.java
@@ -0,0 +1,50 @@
+package io.github.jshipit;
+
+import com.fasterxml.jackson.databind.JsonNode;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+
+public class JshipIT {
+
+    public int finishedCount;
+    public int downloadThreads;
+
+    public JshipIT(String[] args) {
+        DockerAPIHelper api = new DockerAPIHelper("registry.getcryst.al","crystal/misc", "docker", "latest");
+        JsonNode manifest = null;
+
+        System.out.println("API Token: " + api.getApiToken());
+        try {
+            manifest = api.fetchManifestJson();
+        } catch (IOException e) {
+
+        }
+
+        System.out.println("Manifest: " + manifest);
+
+        Path path = Path.of("./tmp_"+api.getImage()+"_"+api.getTag());
+        try {
+            Files.createDirectory(path);
+        } catch (IOException e) {
+            System.out.println("Failed to create directory: " + path);
+            e.printStackTrace();
+            return;
+        }
+
+        finishedCount = 0;
+        downloadThreads = 0;
+
+        JsonNode layers = manifest.get("layers");
+        for (JsonNode layer : layers) {
+            System.out.println("Layer: " + layer);
+            try {
+                api.fetchBlob(layer.get("digest").asText(), path.toString());
+                downloadThreads = downloadThreads+1;
+            } catch (IOException e) {
+                e.printStackTrace();
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/main/java/io/github/jshipit/Main.java b/src/main/java/io/github/jshipit/Main.java
index df1cea7..9de8b6d 100755..100644
--- a/src/main/java/io/github/jshipit/Main.java
+++ b/src/main/java/io/github/jshipit/Main.java
@@ -1,41 +1,9 @@
 package io.github.jshipit;
 
-import com.fasterxml.jackson.databind.JsonNode;
-
-import java.io.IOException;
-import java.nio.file.Files;
-import java.nio.file.Path;
-
 public class Main {
-    public static void main(String[] args) {
-        DockerAPIHelper api = new DockerAPIHelper("registry.docker.io", "https://auth.docker.io/token", "library", "archlinux", "latest");
-        JsonNode manifest = null;
-
-        System.out.println("API Token: " + api.getApiToken());
-        try {
-            manifest = api.fetchManifestJson();
-        } catch (IOException e) {
-
-        }
 
-        System.out.println("Manifest: " + manifest);
-
-        Path path = Path.of("./tmp_"+api.getImage()+"_"+api.getTag());
-        try {
-            Files.createDirectory(path);
-        } catch (IOException e) {
-            System.out.println("Failed to create directory: " + path);
-            return;
-        }
-
-        JsonNode layers = manifest.get("layers");
-        for (JsonNode layer : layers) {
-            System.out.println("Layer: " + layer);
-            try {
-                api.fetchBlob(layer.get("digest").asText(), path.toString());
-            } catch (IOException e) {
-
-            }
-        }
+    public static void main(String[] args) {
+        JshipIT prog = new JshipIT(args); // I HATE STATIC FUNCTIONS
     }
-}
\ No newline at end of file
+
+}