diff --git a/go.mod b/go.mod index 87e1729..f8d5680 100644 --- a/go.mod +++ b/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 ) diff --git a/go.sum b/go.sum index 991d482..5a2e358 100644 --- a/go.sum +++ b/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= diff --git a/pkg/cli/root.go b/pkg/cli/root.go index 17e02b5..44674ee 100644 --- a/pkg/cli/root.go +++ b/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) diff --git a/pkg/cli/size.go b/pkg/cli/size.go new file mode 100644 index 0000000..1f9cdd8 --- /dev/null +++ b/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 +}