pulli.go

package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
)

type GetTokenResponse struct {
	Token string `json:"token"`
}

type GetImageManifestResponse struct {
	Manifests     []*Manifest `json:"manifests"`
	MediaType     string      `json:"mediaType"`
	SchemaVersion int         `json:"schemaVersion"`
}

type Manifest struct {
	Annotations map[string]string `json:"annotations"`
	Digest      string            `json:"digest"`
	MediaType   string            `json:"mediaType"`
	Platform    struct {
		Architecture string `json:"architecture"`
		OS           string `json:"os"`
	} `json:"platform"`
	Size int `json:"size"`
}

type GetSpecificManifestResponse struct {
	SchemaVersion int    `json:"schemaVersion"`
	MediaType     string `json:"mediaType"`
	Config        *struct {
		MediaType string `json:"mediaType"`
		Digest    string `json:"digest"`
		Size      int    `json:"size"`
	}
	Layers []*struct {
		MediaType   string            `json:"mediaType"`
		Digest      string            `json:"digest"`
		Size        int               `json:"size"`
		Annotations map[string]string `json:"annotations"`
	}
	Annotations map[string]string `json:"annotations"`
}

func usage() {
	prog := os.Args[0]
	fmt.Fprintf(os.Stderr, "usage: %s <image-name> [architecture] [os]\n", prog)
	fmt.Fprintf(os.Stderr, "architecture: default to amd64, os: default to linux\n")
	fmt.Fprintf(os.Stderr, "currently only supports pulling from Dockerhub\n")
}

func main() {
	if len(os.Args) < 2 {
		usage()
		os.Exit(-1)
	}

	image := os.Args[1]
	arch := "amd64"
	operatingSystem := "linux"

	if len(os.Args) == 3 {
		arch = os.Args[2]
	}

	if len(os.Args) == 4 {
		arch = os.Args[2]
		operatingSystem = os.Args[3]
	}

	token, err := getToken(image)
	if err != nil {
		fmt.Printf("failed to get bearer token: %s\n", err)
		os.Exit(-1)
	}

	fmt.Fprintf(os.Stderr, "fetching image manifest for %s (arch: %s, os: %s)\n", image, arch, operatingSystem)

	imageManifestResponse, err := getImageManifest(image, token)
	if err != nil {
		fmt.Printf("failed to get image manifests: %s\n", err)
		os.Exit(-1)
	}

	platformSpecificManifests := make([]*Manifest, 0)
	for _, manifest := range imageManifestResponse.Manifests {
		if manifest.Platform.Architecture == arch && manifest.Platform.OS == operatingSystem {
			platformSpecificManifests = append(platformSpecificManifests, manifest)
		}
	}

	if len(platformSpecificManifests) == 0 {
		fmt.Fprintf(os.Stderr, "No image manifest found for %s\n", image)
		os.Exit(-1)
	}

	specificManifest := platformSpecificManifests[0]

	fmt.Fprintf(os.Stderr, "getting image specific manifest for %s (arch=%s os=%s digest=%s)\n", image, arch, operatingSystem, specificManifest.Digest)

	specificImageManifestResponse, err := getSpecificManifest(token, image, specificManifest.Digest)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to get manifest for %s (arch=%s os=%s digest=%s)\n", image, arch, operatingSystem, specificManifest.Digest)
		os.Exit(-1)
	}

	if _, err := isLayerBlobExists(token, image, specificImageManifestResponse.Layers[0].Digest); err != nil {
		fmt.Fprintf(os.Stderr, "%s\n", err)
		os.Exit(-1)
	}

	if err := downloadLayerBlob(token, image, specificImageManifestResponse.Layers[0].Digest); err != nil {
		fmt.Fprintf(os.Stderr, "error downloading image: %s\n", err)
	}
}

func getToken(imageName string) (string, error) {
	endpoint := fmt.Sprintf("https://auth.docker.io/token?service=registry.docker.io&scope=repository:library/%s:pull", imageName)
	resp, err := http.Get(endpoint)
	if err != nil {
		return "", err
	}

	bodyBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("failed to read HTTP response body: %s", err)
	}

	var getTokenResp GetTokenResponse
	if err := json.Unmarshal(bodyBytes, &getTokenResp); err != nil {
		return "", fmt.Errorf("failed to deserialize JSON: %s", err)
	}

	return getTokenResp.Token, nil
}

func getImageManifest(imageName string, token string) (*GetImageManifestResponse, error) {
	endpoint := fmt.Sprintf("https://registry-1.docker.io/v2/library/%s/manifests/latest", imageName)

	req, err := http.NewRequest("GET", endpoint, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", token))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed to get image manifest: %s", err)
	}

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("unexpected HTTP status code %d when getting image manifest", resp.StatusCode)
	}

	respBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read response body: %s", err)
	}

	var getManifestResp GetImageManifestResponse
	if err := json.Unmarshal(respBytes, &getManifestResp); err != nil {
		return nil, fmt.Errorf("failed to deserialize JSON: %s", err)
	}

	return &getManifestResp, nil
}

func getSpecificManifest(token string, imageName string, digest string) (*GetSpecificManifestResponse, error) {
	endpoint := fmt.Sprintf("https://registry-1.docker.io/v2/library/%s/manifests/%s", imageName, digest)
	req, _ := http.NewRequest("GET", endpoint, nil) // TODO: handle error
	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", token))

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}

	respBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	var specificManifestResp GetSpecificManifestResponse
	if err := json.Unmarshal(respBytes, &specificManifestResp); err != nil {
		return nil, err
	}

	return &specificManifestResp, nil
}

func isLayerBlobExists(token string, image string, digest string) (bool, error) {
	endpoint := fmt.Sprintf("https://registry-1.docker.io/v2/library/%s/blobs/%s", image, digest)
	req, err := http.NewRequest("HEAD", endpoint, nil)
	if err != nil {
		return false, fmt.Errorf("failed to create HTTP request: %s", err)
	}

	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", token))

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return false, fmt.Errorf("failed to check layer blob: %s", err)
	}

	if resp.StatusCode != 200 {
		return false, fmt.Errorf("layer blob possibly does not exist. status=%d\n", resp.StatusCode)
	}

	return true, nil
}

func downloadLayerBlob(token string, image string, digest string) error {
	endpoint := fmt.Sprintf("https://registry-1.docker.io/v2/library/%s/blobs/%s", image, digest)
	req, err := http.NewRequest("GET", endpoint, nil)
	if err != nil {
		return err
	}

	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", token))

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("failed to download blob: %s", err)
	}

	r := bufio.NewReader(resp.Body)
	destFile := os.Stdout

	wlen, err := r.WriteTo(destFile)
	if err != nil {
		return fmt.Errorf("write error: %s", err)
	}

	fmt.Fprintf(os.Stderr, "wrote %d bytes\n", wlen)

	return nil
}