feat: upload local disk images (#15)

The new options/flag enables users to use a local file as the image,
instead of a publicly available file from a web server.

cmd/upload.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"fmt"
 	"net/url"
+	"os"
 
 	"github.com/hetznercloud/hcloud-go/v2/hcloud"
 	"github.com/spf13/cobra"
@@ -13,6 +14,7 @@ import (
 
 const (
 	uploadFlagImageURL     = "image-url"
+	uploadFlagImagePath    = "image-path"
 	uploadFlagCompression  = "compression"
 	uploadFlagArchitecture = "architecture"
 	uploadFlagDescription  = "description"
@@ -21,10 +23,14 @@ const (
 
 // uploadCmd represents the upload command
 var uploadCmd = &cobra.Command{
-	Use:   "upload",
+	Use:   "upload (--image-path=<local-path> | --image-url=<url>) --architecture=<x86|arm>",
 	Short: "Upload the specified disk image into your Hetzner Cloud project.",
 	Long: `This command implements a fake "upload", by going through a real server and snapshots.
 This does cost a bit of money for the server.`,
+	Example: `  hcloud-upload-image upload --image-path /home/you/images/custom-linux-image-x86.bz2 --architecture x86 --compression bz2 --description "My super duper custom linux"
+  hcloud-upload-image upload --image-url https://examples.com/image-arm.raw --architecture arm --labels foo=bar,version=latest
+`,
+
 	GroupID: "primary",
 
 	RunE: func(cmd *cobra.Command, args []string) error {
@@ -32,23 +38,36 @@ This does cost a bit of money for the server.`,
 		logger := contextlogger.From(ctx)
 
 		imageURLString, _ := cmd.Flags().GetString(uploadFlagImageURL)
+		imagePathString, _ := cmd.Flags().GetString(uploadFlagImagePath)
 		imageCompression, _ := cmd.Flags().GetString(uploadFlagCompression)
 		architecture, _ := cmd.Flags().GetString(uploadFlagArchitecture)
 		description, _ := cmd.Flags().GetString(uploadFlagDescription)
 		labels, _ := cmd.Flags().GetStringToString(uploadFlagLabels)
 
+		options := hcloudimages.UploadOptions{
+			ImageCompression: hcloudimages.Compression(imageCompression),
+			Architecture:     hcloud.Architecture(architecture),
+			Description:      hcloud.Ptr(description),
+			Labels:           labels,
+		}
+
+		if imageURLString != "" {
 			imageURL, err := url.Parse(imageURLString)
 			if err != nil {
 				return fmt.Errorf("unable to parse url from --%s=%q: %w", uploadFlagImageURL, imageURLString, err)
 			}
 
-		image, err := client.Upload(ctx, hcloudimages.UploadOptions{
+			options.ImageURL = imageURL
-			ImageURL:         imageURL,
+		} else if imagePathString != "" {
-			ImageCompression: hcloudimages.Compression(imageCompression),
+			imageFile, err := os.Open(imagePathString)
-			Architecture:     hcloud.Architecture(architecture),
+			if err != nil {
-			Description:      hcloud.Ptr(description),
+				return fmt.Errorf("unable to read file from --%s=%q: %w", uploadFlagImagePath, imagePathString, err)
-			Labels:           labels,
+			}
-		})
+
+			options.ImageReader = imageFile
+		}
+
+		image, err := client.Upload(ctx, options)
 		if err != nil {
 			return fmt.Errorf("failed to upload the image: %w", err)
 		}
@@ -62,23 +81,25 @@ This does cost a bit of money for the server.`,
 func init() {
 	rootCmd.AddCommand(uploadCmd)
 
-	uploadCmd.Flags().String(uploadFlagImageURL, "", "Remote URL of the disk image that should be uploaded (required)")
+	uploadCmd.Flags().String(uploadFlagImageURL, "", "Remote URL of the disk image that should be uploaded")
-	_ = uploadCmd.MarkFlagRequired(uploadFlagImageURL)
+	uploadCmd.Flags().String(uploadFlagImagePath, "", "Local path to the disk image that should be uploaded")
+	uploadCmd.MarkFlagsMutuallyExclusive(uploadFlagImageURL, uploadFlagImagePath)
+	uploadCmd.MarkFlagsOneRequired(uploadFlagImageURL, uploadFlagImagePath)
 
-	uploadCmd.Flags().String(uploadFlagCompression, "", "Type of compression that was used on the disk image")
+	uploadCmd.Flags().String(uploadFlagCompression, "", "Type of compression that was used on the disk image [choices: bz2]")
 	_ = uploadCmd.RegisterFlagCompletionFunc(
 		uploadFlagCompression,
 		cobra.FixedCompletions([]string{string(hcloudimages.CompressionBZ2)}, cobra.ShellCompDirectiveNoFileComp),
 	)
 
-	uploadCmd.Flags().String(uploadFlagArchitecture, "", "CPU Architecture of the disk image. Choices: x86|arm")
+	uploadCmd.Flags().String(uploadFlagArchitecture, "", "CPU architecture of the disk image [choices: x86, arm]")
 	_ = uploadCmd.RegisterFlagCompletionFunc(
 		uploadFlagArchitecture,
 		cobra.FixedCompletions([]string{string(hcloud.ArchitectureX86), string(hcloud.ArchitectureARM)}, cobra.ShellCompDirectiveNoFileComp),
 	)
 	_ = uploadCmd.MarkFlagRequired(uploadFlagArchitecture)
 
-	uploadCmd.Flags().String(uploadFlagDescription, "", "Description for the resulting Image")
+	uploadCmd.Flags().String(uploadFlagDescription, "", "Description for the resulting image")
 
-	uploadCmd.Flags().StringToString(uploadFlagLabels, map[string]string{}, "Labels for the resulting Image")
+	uploadCmd.Flags().StringToString(uploadFlagLabels, map[string]string{}, "Labels for the resulting image")
 }

hcloudimages/client.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"log/slog"
 	"net/url"
 	"time"
@@ -47,6 +48,10 @@ var (
 type UploadOptions struct {
 	// ImageURL must be publicly available. The instance will download the image from this endpoint.
 	ImageURL *url.URL
+
+	// ImageReader
+	ImageReader io.Reader
+
 	// ImageCompression describes the compression of the referenced image file. It defaults to [CompressionNone]. If
 	// set to anything else, the file will be decompressed before written to the disk.
 	ImageCompression Compression
@@ -207,7 +212,7 @@ func (s *Client) Upload(ctx context.Context, options UploadOptions) (*hcloud.Ima
 	}()
 
 	// 3. Activate Rescue System
-	logger.InfoContext(ctx, "# Step 4: Activating Rescue System")
+	logger.InfoContext(ctx, "# Step 3: Activating Rescue System")
 	enableRescueResult, _, err := s.c.Server.EnableRescue(ctx, server, hcloud.ServerEnableRescueOpts{
 		Type:    defaultRescueType,
 		SSHKeys: []*hcloud.SSHKey{key},
@@ -276,20 +281,24 @@ func (s *Client) Upload(ctx context.Context, options UploadOptions) (*hcloud.Ima
 
 	// 6. SSH On Server: Download Image, Decompress, Write to Root Disk
 	logger.InfoContext(ctx, "# Step 6: Downloading image and writing to disk")
-	decompressionCommand := ""
+	cmd := ""
+	if options.ImageURL != nil {
+		cmd += fmt.Sprintf("wget --no-verbose -O - %q | ", options.ImageURL.String())
+	}
+
 	if options.ImageCompression != CompressionNone {
 		switch options.ImageCompression {
 		case CompressionBZ2:
-			decompressionCommand += "| bzip2 -cd"
+			cmd += "bzip2 -cd | "
 		default:
 			return nil, fmt.Errorf("unknown compression: %q", options.ImageCompression)
 		}
 	}
 
-	fullCmd := fmt.Sprintf("wget --no-verbose -O - %q %s | dd of=/dev/sda bs=4M && sync", options.ImageURL.String(), decompressionCommand)
+	cmd += "dd of=/dev/sda bs=4M && sync"
-	logger.DebugContext(ctx, "running download, decompress and write to disk command", "cmd", fullCmd)
+	logger.DebugContext(ctx, "running download, decompress and write to disk command", "cmd", cmd)
 
-	output, err := sshsession.Run(sshClient, fullCmd)
+	output, err := sshsession.Run(sshClient, cmd, options.ImageReader)
 	logger.InfoContext(ctx, "# Step 6: Finished writing image to disk")
 	logger.DebugContext(ctx, string(output))
 	if err != nil {
@@ -298,7 +307,7 @@ func (s *Client) Upload(ctx context.Context, options UploadOptions) (*hcloud.Ima
 
 	// 7. SSH On Server: Shutdown
 	logger.InfoContext(ctx, "# Step 7: Shutting down server")
-	_, err = sshsession.Run(sshClient, "shutdown now")
+	_, err = sshsession.Run(sshClient, "shutdown now", nil)
 	if err != nil {
 		// TODO Verify if shutdown error, otherwise return
 		logger.WarnContext(ctx, "shutdown returned error", "err", err)

hcloudimages/internal/sshsession/session.go

@@ -1,12 +1,21 @@
 package sshsession
 
-import "golang.org/x/crypto/ssh"
+import (
+	"io"
 
-func Run(client *ssh.Client, cmd string) ([]byte, error) {
+	"golang.org/x/crypto/ssh"
+)
+
+func Run(client *ssh.Client, cmd string, stdin io.Reader) ([]byte, error) {
 	sess, err := client.NewSession()
+
 	if err != nil {
 		return nil, err
 	}
 	defer sess.Close()
+
+	if stdin != nil {
+		sess.Stdin = stdin
+	}
 	return sess.CombinedOutput(cmd)
 }