report on size of a version or all versions

go.mod

@@ -3,8 +3,10 @@ module github.com/anotherjesse/r8im
 go 1.19
 
 require (
+	github.com/dustin/go-humanize v1.0.1
 	github.com/google/go-containerregistry v0.13.0
-	github.com/spf13/cobra v1.6.1
+	github.com/replicate/replicate-go v0.14.2
+	github.com/spf13/cobra v1.7.0
 )
 
 require (
@@ -13,15 +15,15 @@ require (
 	github.com/docker/distribution v2.8.1+incompatible // indirect
 	github.com/docker/docker v20.10.20+incompatible // indirect
 	github.com/docker/docker-credential-helpers v0.7.0 // indirect
-	github.com/inconshreveable/mousetrap v1.0.1 // indirect
-	github.com/klauspost/compress v1.15.11 // indirect
+	github.com/inconshreveable/mousetrap v1.1.0 // indirect
+	github.com/klauspost/compress v1.17.0 // indirect
 	github.com/mitchellh/go-homedir v1.1.0 // indirect
 	github.com/opencontainers/go-digest v1.0.0 // indirect
 	github.com/opencontainers/image-spec v1.1.0-rc2 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
-	github.com/sirupsen/logrus v1.9.0 // indirect
+	github.com/sirupsen/logrus v1.9.3 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/vbatts/tar-split v0.11.2 // indirect
-	golang.org/x/sync v0.1.0 // indirect
-	golang.org/x/sys v0.1.0 // indirect
+	golang.org/x/sync v0.5.0 // indirect
+	golang.org/x/sys v0.14.0 // indirect
 )

go.sum

@@ -14,13 +14,16 @@ github.com/docker/docker v20.10.20+incompatible h1:kH9tx6XO+359d+iAkumyKDc5Q1kOw
 github.com/docker/docker v20.10.20+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
 github.com/docker/docker-credential-helpers v0.7.0 h1:xtCHsjxogADNZcdv1pKUHXryefjlVRqWqIhk/uXJp0A=
 github.com/docker/docker-credential-helpers v0.7.0/go.mod h1:rETQfLdHNT3foU5kuNkFR1R1V12OJRRO5lzt2D1b5X0=
+github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
+github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
 github.com/google/go-containerregistry v0.13.0 h1:y1C7Z3e149OJbOPDBxLYR8ITPz8dTKqQwjErKVHJC8k=
 github.com/google/go-containerregistry v0.13.0/go.mod h1:J9FQ+eSS4a1aC2GNZxvNpbWhgp0487v+cgiilB4FqDo=
-github.com/inconshreveable/mousetrap v1.0.1 h1:U3uMjPSQEBMNp1lFxmllqCPM6P5u/Xq7Pgzkat/bFNc=
-github.com/inconshreveable/mousetrap v1.0.1/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
-github.com/klauspost/compress v1.15.11 h1:Lcadnb3RKGin4FYM/orgq0qde+nc15E5Cbqg4B9Sx9c=
+github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
+github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
 github.com/klauspost/compress v1.15.11/go.mod h1:QPwzmACJjUTFsnSHH934V6woptycfrDDJnH7hvFVbGM=
+github.com/klauspost/compress v1.17.0 h1:Rnbp4K9EjcDuVuHtd0dgA4qNuv9yKDYKK1ulpJwgrqM=
+github.com/klauspost/compress v1.17.0/go.mod h1:ntbaceVETuRiXiv4DpjP66DpAtAGkEQskQzEyD//IeE=
 github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
 github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
 github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
@@ -31,30 +34,32 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/replicate/replicate-go v0.14.2 h1:XgK+REvYrWs7qDeyugxHA93h31qBhEFk/3p1/p2w3W8=
+github.com/replicate/replicate-go v0.14.2/go.mod h1:otIrl1vDmyjNhTzmVmp/mQU3Wt1+3387gFNEsAZq0ig=
 github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/shurcooL/sanitized_anchor_name v1.0.0/go.mod h1:1NzhyTcUVG4SuEtjjoZeVRXNmyL/1OwPU0+IJeTBvfc=
 github.com/sirupsen/logrus v1.7.0/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
-github.com/sirupsen/logrus v1.9.0 h1:trlNQbNUG3OdDrDil03MCb1H2o9nJ1x4/5LYw7byDE0=
-github.com/sirupsen/logrus v1.9.0/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
-github.com/spf13/cobra v1.6.1 h1:o94oiPyS4KD1mPy2fmcYYHHfCxLqYjJOhGsCHFZtEzA=
-github.com/spf13/cobra v1.6.1/go.mod h1:IOw/AERYS7UzyrGinqmz6HLUo219MORXGxhbaJUqzrY=
+github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
+github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
+github.com/spf13/cobra v1.7.0 h1:hyqWnYt1ZQShIddO5kBpj3vu05/++x6tJ6dg8EC572I=
+github.com/spf13/cobra v1.7.0/go.mod h1:uLxZILRyS/50WlhOIKD7W6V5bgeIt+4sICxh6uRMrb0=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
-github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
 github.com/urfave/cli v1.22.4/go.mod h1:Gos4lmkARVdJ6EkW0WaNv/tZAAMe9V7XWyB60NtXRu0=
 github.com/vbatts/tar-split v0.11.2 h1:Via6XqJr0hceW4wff3QRzD5gAk/tatMw/4ZA7cTlIME=
 github.com/vbatts/tar-split v0.11.2/go.mod h1:vV3ZuO2yWSVsz+pfFzDG/upWH1JhjOiEaWq6kXyQ3VI=
 golang.org/x/sync v0.0.0-20201207232520-09787c993a3a/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
-golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.5.0 h1:60k92dhOjHxJkrqnwsfl8KuaHbn/5dl0lUPUklKo3qE=
+golang.org/x/sync v0.5.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.1.0 h1:kunALQeHf1/185U1i0GOB/fy1IPRDDpuoOOqRReG57U=
-golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.14.0 h1:Vz7Qs629MkJkGyHxUlRHizWJRG2j8fbQKjELVSNhy7Q=
+golang.org/x/sys v0.14.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

pkg/cli/root.go

@@ -19,9 +19,10 @@ func NewRootCommand() (*cobra.Command, error) {
 	rootCmd.AddCommand(
 		newAffixCommand(),
 		newCloneCommand(),
-		newLayerCommand(),
 		newExtractCommand(),
+		newLayerCommand(),
 		newRemixCommand(),
+		newSizeCommand(),
 		newZstdCommand(),
 	)
 	logs.Warn = log.New(os.Stderr, "gcr WARN: ", log.LstdFlags)

pkg/cli/size.go

@@ -0,0 +1,153 @@
+package cli
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"strings"
+
+	"github.com/dustin/go-humanize"
+	"github.com/google/go-containerregistry/pkg/authn"
+	"github.com/replicate/replicate-go"
+	"github.com/spf13/cobra"
+
+	"github.com/anotherjesse/r8im/pkg/auth"
+	"github.com/anotherjesse/r8im/pkg/images"
+)
+
+func newSizeCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:    "size",
+		Short:  "calculate size of existing image",
+		Hidden: false,
+		Args:   cobra.ExactArgs(1),
+		RunE:   sizeCommand,
+	}
+
+	cmd.Flags().StringVarP(&sToken, "token", "t", "", "replicate cog token")
+	cmd.Flags().StringVarP(&sRegistry, "registry", "r", "r8.im", "registry host")
+
+	return cmd
+}
+
+func sizeCommand(cmd *cobra.Command, args []string) error {
+	if sToken == "" {
+		sToken = os.Getenv("REPLICATE_API_TOKEN")
+	}
+	if sToken == "" {
+		sToken = os.Getenv("COG_TOKEN")
+	}
+
+	u, err := auth.VerifyCogToken(sRegistry, sToken)
+	if err != nil {
+		fmt.Fprintln(os.Stderr, "authentication error, invalid token or registry host error")
+		return err
+	}
+	auth := authn.FromConfig(authn.AuthConfig{Username: u, Password: sToken})
+
+	all_images, err := getR8Urls(args[0], auth)
+	if err != nil {
+		return err
+	}
+
+	uniqLayers := make(map[string]int64)
+	totalSize := 0
+
+	for _, image := range all_images {
+		imageSize := 0
+		imageUniqSize := 0
+		layers, err := images.Layers(image, auth)
+		if err != nil {
+			return err
+		}
+
+		for _, layer := range layers {
+			imageSize += int(layer.Size)
+			if _, ok := uniqLayers[layer.Digest]; ok {
+				continue
+			}
+			uniqLayers[layer.Digest] = layer.Size
+			imageUniqSize += int(layer.Size)
+		}
+
+		totalSize += imageUniqSize
+		version := strings.Split(image, "@sha256:")[1]
+		fmt.Printf("%s\t%s\t%s\n", version, humanize.Bytes(uint64(imageSize)), humanize.Bytes(uint64(imageUniqSize)))
+	}
+
+	fmt.Printf("Total Size: %s\n", humanize.Bytes(uint64(totalSize)))
+
+	return nil
+}
+
+func ensureRegistry(baseRef string) string {
+	if !strings.Contains(baseRef, sRegistry) {
+		return sRegistry + "/" + baseRef
+	}
+	return baseRef
+}
+
+func getReplicateClient(session authn.Authenticator) (*replicate.Client, error) {
+	a, err := session.Authorization()
+	if err != nil {
+		return nil, err
+	}
+	token := a.Password
+	client, err := replicate.NewClient(replicate.WithToken(token))
+	if err != nil {
+		return nil, err
+	}
+	return client, nil
+}
+
+func getVersions(owner string, model string, session authn.Authenticator) ([]string, error) {
+	client, err := getReplicateClient(session)
+	if err != nil {
+		return nil, err
+	}
+
+	resp, err := client.ListModelVersions(context.Background(), owner, model)
+	if err != nil {
+		return nil, err
+	}
+
+	var versions []string
+	for _, version := range resp.Results {
+		versions = append(versions, version.ID)
+	}
+
+	return versions, nil
+}
+
+func getR8Urls(baseRef string, session authn.Authenticator) ([]string, error) {
+	baseRef = ensureRegistry(baseRef)
+	parts := strings.Split(baseRef, "/")
+	owner := parts[len(parts)-2]
+	model := parts[len(parts)-1]
+	version := ""
+	if strings.Contains(model, ":") {
+		version = strings.Split(model, ":")[1]
+		model = strings.Split(model, ":")[0]
+	}
+	if strings.Contains(model, "@") {
+		model = strings.Split(model, "@")[0]
+	}
+
+	if version != "" {
+		image := fmt.Sprintf("%s/%s/%s@sha256:%s", sRegistry, owner, model, version)
+		return []string{image}, nil
+	}
+
+	versions, err := getVersions(owner, model, session)
+	if err != nil {
+		return nil, err
+	}
+
+	var all_images []string
+	for _, version := range versions {
+		image := fmt.Sprintf("%s/%s/%s@sha256:%s", sRegistry, owner, model, version)
+		all_images = append(all_images, image)
+	}
+
+	return all_images, nil
+}