From 24b2123582b047aa189cafa7c9322dd717425bb6 Mon Sep 17 00:00:00 2001
From: Mikhail Klementev <blame@dumpstack.io>
Date: Mon, 19 Aug 2019 23:02:34 +0000
Subject: [PATCH] More flexible way to change smep/smap/kaslr while debug

---
 debug.go | 27 ++++++++++++++++++---------
 main.go  | 25 +++++++++++++++++++++----
 2 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/debug.go b/debug.go
index f4bc16c..1375d1b 100644
--- a/debug.go
+++ b/debug.go
@@ -83,7 +83,8 @@ func interactive(q *qemu.System) (err error) {
 }
 
 func debugHandler(kcfg config.KernelConfig, workPath, kernRegex, gdb string,
-	dockerTimeout time.Duration, kaslr, smep, smap bool) (err error) {
+	dockerTimeout time.Duration, yekaslr, yesmep, yesmap,
+	nokaslr, nosmep, nosmap bool) (err error) {
 
 	ka, err := config.ReadArtifactConfig(workPath + "/.out-of-tree.toml")
 	if err != nil {
@@ -115,18 +116,26 @@ func debugHandler(kcfg config.KernelConfig, workPath, kernRegex, gdb string,
 	fmt.Printf("[*] SMP: %d CPUs\n", q.Cpus)
 	fmt.Printf("[*] Memory: %d MB\n", q.Memory)
 
-	q.SetKASLR(kaslr)
+	q.SetKASLR(false) // set KASLR to false by default because of gdb
+	q.SetSMEP(!ka.Mitigations.DisableSmep)
+	q.SetSMAP(!ka.Mitigations.DisableSmap)
 
-	if !smep {
-		q.SetSMEP(false)
-	} else {
-		q.SetSMEP(!ka.Mitigations.DisableSmep)
+	if yekaslr {
+		q.SetKASLR(true)
+	} else if nokaslr {
+		q.SetKASLR(false)
 	}
 
-	if !smap {
+	if yesmep {
+		q.SetSMEP(true)
+	} else if nosmep {
+		q.SetSMEP(false)
+	}
+
+	if yesmap {
+		q.SetSMAP(true)
+	} else if nosmap {
 		q.SetSMAP(false)
-	} else {
-		q.SetSMAP(!ka.Mitigations.DisableSmap)
 	}
 
 	redgreen := func(name string, enabled bool) aurora.Value {
diff --git a/main.go b/main.go
index 094bd1f..ffd4001 100644
--- a/main.go
+++ b/main.go
@@ -177,9 +177,13 @@ func main() {
 	debugFlagGDB := debugCommand.Flag("gdb", "Set gdb listen address")
 	debugGDB := debugFlagGDB.Default("tcp::1234").String()
 
-	kaslr := debugCommand.Flag("enable-kaslr", "Enable KASLR").Default("false").Bool()
-	nosmep := debugCommand.Flag("disable-smep", "Disable SMEP").Default("false").Bool()
-	nosmap := debugCommand.Flag("disable-smap", "Disable SMAP").Default("false").Bool()
+	yekaslr := debugCommand.Flag("enable-kaslr", "Enable KASLR").Bool()
+	yesmep := debugCommand.Flag("enable-smep", "Enable SMEP").Bool()
+	yesmap := debugCommand.Flag("enable-smap", "Enable SMAP").Bool()
+
+	nokaslr := debugCommand.Flag("disable-kaslr", "Disable KASLR").Bool()
+	nosmep := debugCommand.Flag("disable-smep", "Disable SMEP").Bool()
+	nosmap := debugCommand.Flag("disable-smap", "Disable SMAP").Bool()
 
 	bootstrapCommand := app.Command("bootstrap",
 		"Create directories && download images")
@@ -238,6 +242,18 @@ func main() {
 
 	kingpin.MustParse(app.Parse(os.Args[1:]))
 
+	if *yekaslr && *nokaslr {
+		log.Fatalln("Only one of disable/enable can be used at once")
+	}
+
+	if *yesmep && *nosmep {
+		log.Fatalln("Only one of disable/enable can be used at once")
+	}
+
+	if *yesmap && *nosmap {
+		log.Fatalln("Only one of disable/enable can be used at once")
+	}
+
 	kcfg, err := config.ReadKernelConfig(*kcfgPath)
 	if err != nil {
 		log.Println(err)
@@ -283,7 +299,8 @@ func main() {
 		err = genConfig(config.KernelExploit)
 	case debugCommand.FullCommand():
 		err = debugHandler(kcfg, *path, *debugKernel, *debugGDB,
-			*dockerTimeout, *kaslr, !*nosmep, !*nosmap)
+			*dockerTimeout, *yekaslr, *yesmep, *yesmap,
+			*nokaslr, *nosmep, *nosmap)
 	case bootstrapCommand.FullCommand():
 		err = bootstrapHandler()
 	case logQueryCommand.FullCommand():