From 5dbbb332970a486cfc7768aa75607a98192ab64c Mon Sep 17 00:00:00 2001
From: Mikhail Klementev <blame@dumpstack.io>
Date: Mon, 12 Aug 2019 22:58:34 +0000
Subject: [PATCH] Implements parameter "--max=X" for autogen

---
 kernel.go | 39 ++++++++++++++++++++++++++++++++-------
 main.go   | 11 ++++++++++-
 pew.go    |  3 +++
 3 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/kernel.go b/kernel.go
index 309a4e1..6fb09df 100644
--- a/kernel.go
+++ b/kernel.go
@@ -9,6 +9,8 @@ import (
 	"fmt"
 	"io/ioutil"
 	"log"
+	"math"
+	"math/rand"
 	"os"
 	"os/exec"
 	"os/user"
@@ -20,6 +22,8 @@ import (
 	"github.com/naoina/toml"
 )
 
+var KERNELS_ALL int64 = math.MaxInt64
+
 func kernelListHandler(kcfg config.KernelConfig) (err error) {
 	if len(kcfg.Kernels) == 0 {
 		return errors.New("No kernels found")
@@ -393,7 +397,17 @@ func hasKernel(ki config.KernelInfo, kcfg config.KernelConfig) bool {
 	return false
 }
 
-func generateKernels(km config.KernelMask) (err error) {
+func shuffle(a []string) []string {
+	// Fisher–Yates shuffle
+	for i := len(a) - 1; i > 0; i-- {
+		j := rand.Intn(i + 1)
+		a[i], a[j] = a[j], a[i]
+	}
+	return a
+}
+
+func generateKernels(km config.KernelMask, max int64) (err error) {
+	log.Println("Generating for kernel mask", km)
 	err = generateBaseDockerImage(km)
 	if err != nil {
 		return
@@ -406,9 +420,20 @@ func generateKernels(km config.KernelMask) (err error) {
 		return
 	}
 
-	for i, pkg := range pkgs {
-		log.Println(i, "/", len(pkgs))
-		dockerImageAppend(km, pkg)
+	for i, pkg := range shuffle(pkgs) {
+		if max <= 0 {
+			log.Println("Max is reached")
+			break
+		}
+
+		log.Println(i, "/", len(pkgs), pkg)
+
+		err = dockerImageAppend(km, pkg)
+		if err == nil {
+			max -= 1
+		} else {
+			log.Println("dockerImageAppend", err)
+		}
 	}
 
 	err = kickImage(km.DockerName())
@@ -425,7 +450,7 @@ func generateKernels(km config.KernelMask) (err error) {
 	return
 }
 
-func kernelAutogenHandler(workPath string) (err error) {
+func kernelAutogenHandler(workPath string, max int64) (err error) {
 	ka, err := config.ReadArtifactConfig(workPath + "/.out-of-tree.toml")
 	if err != nil {
 		return
@@ -437,7 +462,7 @@ func kernelAutogenHandler(workPath string) (err error) {
 			return
 		}
 
-		err = generateKernels(sk)
+		err = generateKernels(sk, max)
 		if err != nil {
 			return
 		}
@@ -499,7 +524,7 @@ func kernelGenallHandler(distro, version string) (err error) {
 		DistroRelease: version,
 		ReleaseMask:   ".*",
 	}
-	err = generateKernels(km)
+	err = generateKernels(km, KERNELS_ALL)
 	if err != nil {
 		return
 	}
diff --git a/main.go b/main.go
index 0af4fdd..6fb7ac4 100644
--- a/main.go
+++ b/main.go
@@ -7,10 +7,12 @@ package main
 import (
 	"fmt"
 	"log"
+	"math/rand"
 	"os"
 	"os/exec"
 	"os/user"
 	"sort"
+	"time"
 
 	kingpin "gopkg.in/alecthomas/kingpin.v2"
 
@@ -70,6 +72,8 @@ func checkDockerPermissions() (err error) {
 }
 
 func main() {
+	rand.Seed(time.Now().UnixNano())
+
 	app := kingpin.New(
 		"out-of-tree",
 		"kernel {module, exploit} development tool",
@@ -117,6 +121,11 @@ func main() {
 	kernelListCommand := kernelCommand.Command("list", "List kernels")
 	kernelAutogenCommand := kernelCommand.Command("autogen",
 		"Generate kernels based on a current config")
+	kernelAutogenMax := kernelAutogenCommand.Flag("max",
+		"Download random kernels from set defined by regex in "+
+			"release_mask, but no more than X for each of "+
+			"release_mask").PlaceHolder("X").Default(
+		fmt.Sprint(KERNELS_ALL)).Int64()
 	kernelDockerRegenCommand := kernelCommand.Command("docker-regen",
 		"Regenerate kernels config from out_of_tree_* docker images")
 	kernelGenallCommand := kernelCommand.Command("genall",
@@ -200,7 +209,7 @@ func main() {
 	case kernelListCommand.FullCommand():
 		err = kernelListHandler(kcfg)
 	case kernelAutogenCommand.FullCommand():
-		err = kernelAutogenHandler(*path)
+		err = kernelAutogenHandler(*path, *kernelAutogenMax)
 	case kernelDockerRegenCommand.FullCommand():
 		err = kernelDockerRegenHandler()
 	case kernelGenallCommand.FullCommand():
diff --git a/pew.go b/pew.go
index 557da27..c2affe9 100644
--- a/pew.go
+++ b/pew.go
@@ -39,6 +39,9 @@ func dockerRun(timeout time.Duration, container, workdir, command string) (
 
 	raw, err := cmd.CombinedOutput()
 	if err != nil {
+		e := fmt.Sprintf("error `%v` for cmd `%v` with output `%v`",
+			err, command, string(raw))
+		err = errors.New(e)
 		return
 	}