From aa2ca58c8dd42136de5f627b821d0f739d5399b5 Mon Sep 17 00:00:00 2001
From: Lars Karlslund <lars@karlslund.dk>
Date: Mon, 30 May 2022 11:30:55 +0200
Subject: [PATCH] Used PwnMethodBitmap functionality rather than duplicating it

---
 modules/engine/analyzeobjects.go | 18 +++++-------------
 modules/engine/pwn.go            | 11 ++++++++---
 2 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/modules/engine/analyzeobjects.go b/modules/engine/analyzeobjects.go
index 9762567..f123b97 100644
--- a/modules/engine/analyzeobjects.go
+++ b/modules/engine/analyzeobjects.go
@@ -159,19 +159,11 @@ func AnalyzeObjects(opts AnalyzeObjectsOptions) (pg PwnGraph) {
 					}
 				}
 
-				maxprobability := Probability(-128)
-				for i := 0; i < len(pwnnums); i++ {
-					if detectedmethods.IsSet(PwnMethod(i)) {
-						var probability Probability
-						if forward {
-							probability = CalculateProbability(pwntarget, object, PwnMethod(i))
-						} else {
-							probability = CalculateProbability(object, pwntarget, PwnMethod(i))
-						}
-						if probability > maxprobability {
-							maxprobability = probability
-						}
-					}
+				var maxprobability Probability
+				if forward {
+					maxprobability = detectedmethods.MaxProbability(pwntarget, object)
+				} else {
+					maxprobability = detectedmethods.MaxProbability(object, pwntarget)
 				}
 				if maxprobability < Probability(opts.MinProbability) {
 					// Too unlikeliy, so we skip it
diff --git a/modules/engine/pwn.go b/modules/engine/pwn.go
index 64ea978..b2a6b44 100644
--- a/modules/engine/pwn.go
+++ b/modules/engine/pwn.go
@@ -20,6 +20,11 @@ const MAXPWNMETHODPOSSIBLE = PMBSIZE * 64
 type PwnMethodBitmap [PMBSIZE]uint64
 type Probability int8
 
+const (
+	MINPROBABILITY Probability = -1
+	MAXPROBABILITY Probability = 100
+)
+
 type PwnInfo struct {
 	Target      *Object
 	Method      PwnMethod
@@ -220,12 +225,12 @@ func (m PwnMethodBitmap) IsSet(method PwnMethod) bool {
 	return (m[method/64] & (1 << (method % 64))) != 0 // Uuuuh, nasty and unreadable
 }
 
-func (m PwnMethodBitmap) MaxProbabiltity(source, target *Object) Probability {
+func (m PwnMethodBitmap) MaxProbability(source, target *Object) Probability {
 	var max Probability
 	for i := 0; i < len(pwnnums); i++ {
 		if m.IsSet(PwnMethod(i)) {
-			prob := CalculateProbability(source, target, PwnMethod(i))
-			if prob == 100 {
+			prob := PwnMethod(i).Probability(source, target)
+			if prob == MAXPROBABILITY {
 				return prob
 			}
 			if prob > max {