From 3ec919abc1ce51c1eb908726cc1e130276bc372e Mon Sep 17 00:00:00 2001
From: Mikhail Klementev <blame@dumpstack.io>
Date: Wed, 15 Feb 2023 10:17:57 +0000
Subject: [PATCH] Image editing command

---
 images.go | 93 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 main.go   |  1 +
 2 files changed, 94 insertions(+)

diff --git a/images.go b/images.go
index a8de113..53df892 100644
--- a/images.go
+++ b/images.go
@@ -5,14 +5,107 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
 	"net/http"
 	"os"
 	"os/exec"
+	"os/user"
+	"time"
+
+	"code.dumpstack.io/tools/out-of-tree/config"
+	"code.dumpstack.io/tools/out-of-tree/qemu"
 )
 
+type ImageCmd struct {
+	List ImageListCmd `cmd:"" help:"list images"`
+	Edit ImageEditCmd `cmd:"" help:"edit image"`
+}
+
+type ImageListCmd struct{}
+
+func (cmd *ImageListCmd) Run(g *Globals) (err error) {
+	usr, err := user.Current()
+	if err != nil {
+		return
+	}
+
+	entries, err := os.ReadDir(usr.HomeDir + "/.out-of-tree/images/")
+	if err != nil {
+		return
+	}
+
+	for _, e := range entries {
+		fmt.Println(e.Name())
+	}
+
+	return
+}
+
+type ImageEditCmd struct {
+	Name string `help:"image name" required:""`
+}
+
+func (cmd *ImageEditCmd) Run(g *Globals) (err error) {
+	usr, err := user.Current()
+	if err != nil {
+		return
+	}
+
+	image := usr.HomeDir + "/.out-of-tree/images/" + cmd.Name
+	if !exists(image) {
+		fmt.Println("image does not exist")
+	}
+
+	kcfg, err := config.ReadKernelConfig(g.Config.Kernels)
+	if err != nil {
+		return
+	}
+
+	if len(kcfg.Kernels) == 0 {
+		return errors.New("No kernels found")
+	}
+
+	ki := config.KernelInfo{}
+	for _, k := range kcfg.Kernels {
+		if k.RootFS == image {
+			ki = k
+			break
+		}
+	}
+
+	kernel := qemu.Kernel{
+		KernelPath: ki.KernelPath,
+		InitrdPath: ki.InitrdPath,
+	}
+
+	q, err := qemu.NewSystem(qemu.X86x64, kernel, ki.RootFS)
+
+	q.Mutable = true
+
+	err = q.Start()
+	if err != nil {
+		fmt.Println("Qemu start error:", err)
+		return
+	}
+	defer q.Stop()
+
+	fmt.Print("ssh command:\n\n\t")
+	fmt.Println(q.GetSSHCommand())
+
+	fmt.Print("\npress enter to stop")
+	fmt.Scanln()
+
+	q.Command("root", "poweroff")
+
+	for !q.Died {
+		time.Sleep(time.Second)
+	}
+	return
+}
+
 // inspired by Edd Turtle code
 func downloadFile(filepath string, url string) (err error) {
 	out, err := os.Create(filepath)
diff --git a/main.go b/main.go
index 37e02ab..4a4b512 100644
--- a/main.go
+++ b/main.go
@@ -30,6 +30,7 @@ type CLI struct {
 	Log    LogCmd    `cmd:"" help:"query logs"`
 	Pack   PackCmd   `cmd:"" help:"exploit pack test"`
 	Gen    GenCmd    `cmd:"" help:"generate .out-of-tree.toml skeleton"`
+	Image  ImageCmd  `cmd:"" help:"manage images"`
 
 	Version VersionFlag `name:"version" help:"print version information and quit"`
 }