Wrapped map in object Edge methods to be able to do locking

lkarlslund · Sep 12, 2022 · ecaa968 · ecaa968
1 parent a14a7d3
commit ecaa968
Show file tree

Hide file tree

Showing 8 changed files with 135 additions and 74 deletions.
diff --git a/modules/analyze/webservicefuncs.go b/modules/analyze/webservicefuncs.go
@@ -865,7 +865,7 @@ func analysisfuncs(ws *webservice) {
 
 		var pwnlinks int
 		for _, object := range ws.Objs.Slice() {
-			pwnlinks += len(object.CanPwn)
+			pwnlinks += object.EdgeCount(engine.Out)
 		}
 		result.Statistics["Total"] = len(ws.Objs.Slice())
 		result.Statistics["PwnConnections"] = pwnlinks

diff --git a/modules/engine/analyzeobjects.go b/modules/engine/analyzeobjects.go
@@ -135,31 +135,28 @@ func AnalyzeObjects(opts AnalyzeObjectsOptions) (pg Graph) {
 
 			newconnectionsmap := make(map[ObjectPair]EdgeBitmap) // Pwn Connection between objects
 
-			var ec EdgeConnections
+			var ec EdgeDirection
 			if forward {
-				ec = object.PwnableBy
+				ec = In
 			} else {
-				ec = object.CanPwn
+				ec = Out
 			}
 
 			// Iterate over ever outgoing pwn
 			// This is not efficient, but we sort the pwnlist first
-			for _, target := range ec.Objects() {
-				eb := ec[target]
-
+			object.EdgeIterator(ec, func(target *Object, eb EdgeBitmap) bool {
 				// If this is not a chosen method, skip it
 				detectededges := eb.Intersect(detectedges)
 
-				edgecount := detectededges.Count()
-				if edgecount == 0 {
+				if detectededges.IsBlank() {
 					// Nothing useful or just a deny ACL, skip it
-					continue
+					return true // continue
 				}
 
 				if detectobjecttypes != nil {
 					if _, found := detectobjecttypes[target.Type()]; !found {
 						// We're filtering on types, and it's not wanted
-						continue
+						return true //continue
 					}
 				}
 
@@ -171,7 +168,7 @@ func AnalyzeObjects(opts AnalyzeObjectsOptions) (pg Graph) {
 				}
 				if maxprobability < Probability(opts.MinProbability) {
 					// Too unlikeliy, so we skip it
-					continue
+					return true // continue
 				}
 
 				// If we allow backlinks, all pwns are mapped, no matter who is the victim
@@ -193,19 +190,21 @@ func AnalyzeObjects(opts AnalyzeObjectsOptions) (pg Graph) {
 					// If SIDs match between objects, it's a cross forest link and we want to see it
 					(object.SID().IsNull() || target.SID().IsNull() || object.SID().Component(2) != 21 || object.SID() != target.SID()) {
 					// skip it
-					continue
+					return true // continue
 				}
 
 				if opts.ExcludeObjects != nil {
 					if _, found := opts.ExcludeObjects.FindByID(target.ID()); found {
 						// skip excluded objects
 						// ui.Debug().Msgf("Excluding target %v", pwntarget.DN())
-						continue
+						return true // continue
 					}
 				}
 
 				newconnectionsmap[ObjectPair{Source: object, Target: target}] = detectededges
-			}
+
+				return true
+			})
 
 			if opts.MaxOutgoingConnections == -1 || len(newconnectionsmap) < opts.MaxOutgoingConnections {
 				for pwnpair, detectedmethods := range newconnectionsmap {

diff --git a/modules/engine/analyzepaths.go b/modules/engine/analyzepaths.go
@@ -166,22 +166,20 @@ func AnalyzePaths(start, end *Object, obs *Objects, lookforedges EdgeBitmap, min
 
 		visited[source] = struct{}{}
 
-		for target, edges := range source.CanPwn {
+		source.EdgeIterator(Out, func(target *Object, edges EdgeBitmap) bool {
 			if _, found := visited[target]; !found {
-
 				// If this is not a chosen method, skip it
 				detectededges := edges.Intersect(lookforedges)
 
-				edgecount := detectededges.Count()
-				if edgecount == 0 {
+				if detectededges.IsBlank() {
 					// Nothing useful or just a deny ACL, skip it
-					continue
+					return true //continue
 				}
 
 				prob := detectededges.MaxProbability(v.Object, target)
 				if prob < minprobability {
 					// Skip entirely if too
-					continue
+					return true //continue
 				}
 
 				weight := uint32(101 - prob)
@@ -201,7 +199,8 @@ func AnalyzePaths(start, end *Object, obs *Objects, lookforedges EdgeBitmap, min
 					q.Push(target, sdist+weight)
 				}
 			}
-		}
+			return true
+		})
 	}
 
 	if prev[end] == nil {
@@ -228,7 +227,7 @@ func AnalyzePaths(start, end *Object, obs *Objects, lookforedges EdgeBitmap, min
 			GraphEdge{
 				Source:     prenode,
 				Target:     curnode,
-				EdgeBitmap: prenode.CanPwn[curnode],
+				EdgeBitmap: prenode.Edge(Out, curnode),
 			})
 		if prenode == start {
 			break

diff --git a/modules/engine/edge.go b/modules/engine/edge.go
@@ -74,6 +74,15 @@ func (eb EdgeBitmap) Count() int {
 	return ones
 }
 
+func (eb EdgeBitmap) IsBlank() bool {
+	for i := 0; i < PMBSIZE; i++ {
+		if eb[i] != 0 {
+			return false
+		}
+	}
+	return true
+}
+
 func (eb EdgeBitmap) Edges() []Edge {
 	result := make([]Edge, eb.Count())
 	var n int
@@ -239,6 +248,13 @@ type PwnMethodsAndProbabilities struct {
 }
 */
 
+type EdgeDirection int
+
+const (
+	Out EdgeDirection = 0
+	In  EdgeDirection = 1
+)
+
 type EdgeConnections map[*Object]EdgeBitmap //sAndProbabilities
 
 var globalEdgeConnectionsLock sync.Mutex // Ugly but it will do

diff --git a/modules/engine/object.go b/modules/engine/object.go
@@ -48,13 +48,12 @@ func setThreadsafe(enable bool) {
 var UnknownGUID = uuid.UUID{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}
 
 type Object struct {
-	values    AttributeValueMap
-	PwnableBy EdgeConnections
-	CanPwn    EdgeConnections
-	parent    *Object
-	sdcache   *SecurityDescriptor
-	sid       windowssecurity.SID
-	children  []*Object
+	values   AttributeValueMap
+	edges    [2]EdgeConnections
+	parent   *Object
+	sdcache  *SecurityDescriptor
+	sid      windowssecurity.SID
+	children []*Object
 
 	members          map[*Object]struct{}
 	membersrecursive map[*Object]struct{}
@@ -190,20 +189,20 @@ func (o *Object) AbsorbEx(source *Object, fast bool) {
 		}
 	}
 
-	for pwntarget, methods := range source.CanPwn {
-		target.CanPwn[pwntarget] = target.CanPwn[pwntarget].Merge(methods)
-		delete(source.CanPwn, pwntarget)
+	for pwntarget, methods := range source.edges[Out] {
+		target.edges[Out][pwntarget] = target.edges[Out][pwntarget].Merge(methods)
+		delete(source.edges[Out], pwntarget)
 
-		pwntarget.PwnableBy[target] = pwntarget.PwnableBy[target].Merge(methods)
-		delete(pwntarget.PwnableBy, source)
+		pwntarget.edges[In][target] = pwntarget.edges[In][target].Merge(methods)
+		delete(pwntarget.edges[In], source)
 	}
 
-	for pwner, methods := range source.PwnableBy {
-		target.PwnableBy[pwner] = target.PwnableBy[pwner].Merge(methods)
-		delete(source.PwnableBy, pwner)
+	for pwner, methods := range source.edges[In] {
+		target.edges[In][pwner] = target.edges[In][pwner].Merge(methods)
+		delete(source.edges[In], pwner)
 
-		pwner.CanPwn[target] = pwner.CanPwn[target].Merge(methods)
-		delete(pwner.CanPwn, source)
+		pwner.edges[Out][target] = pwner.edges[Out][target].Merge(methods)
+		delete(pwner.edges[Out], source)
 	}
 
 	// For everyone that is a member of source, relink them to the target
@@ -981,9 +980,9 @@ func (o *Object) init() {
 	if o.values == nil {
 		o.values = NewAttributeValueMap()
 	}
-	if o.CanPwn == nil || o.PwnableBy == nil {
-		o.CanPwn = make(EdgeConnections)
-		o.PwnableBy = make(EdgeConnections)
+	if o.edges[Out] == nil || o.edges[In] == nil {
+		o.edges[Out] = make(EdgeConnections)
+		o.edges[In] = make(EdgeConnections)
 	}
 }
 
@@ -1107,13 +1106,21 @@ func (o *Object) GUID() uuid.UUID {
 	return guid
 }
 
+// Look up edge
+func (o *Object) Edge(direction EdgeDirection, target *Object) EdgeBitmap {
+	o.lock()
+	bm := o.edges[direction][target]
+	o.unlock()
+	return bm
+}
+
 // Register that this object can pwn another object using the given method
 func (o *Object) EdgeTo(target *Object, method Edge) {
-	o.PwnsEx(target, method, false)
+	o.EdgeToEx(target, method, false)
 }
 
 // Enhanched Pwns function that allows us to force the pwn (normally self-pwns are filtered out)
-func (o *Object) PwnsEx(target *Object, method Edge, force bool) {
+func (o *Object) EdgeToEx(target *Object, method Edge, force bool) {
 	if !force {
 		if o == target { // SID check solves (some) dual-AD analysis problems
 			// We don't care about self owns
@@ -1134,14 +1141,49 @@ func (o *Object) PwnsEx(target *Object, method Edge, force bool) {
 	}
 
 	o.lock()
-	o.CanPwn.Set(target, method) // Add the connection
+	o.edges[Out].Set(target, method) // Add the connection
+	o.unlock()
+
+	target.lock()
+	target.edges[In].Set(o, method) // Add the reverse connection too
+	target.unlock()
+}
+
+func (o *Object) EdgeClear(target *Object, method Edge) {
+	o.lock()
+	currentedge := o.edges[Out][target]
+	if currentedge.IsSet(method) {
+		o.edges[Out][target] = currentedge.Clear(method)
+	}
 	o.unlock()
 
 	target.lock()
-	target.PwnableBy.Set(o, method) // Add the reverse connection too
+	currentedge = target.edges[In][o]
+	if currentedge.IsSet(method) {
+		target.edges[In][o] = currentedge.Clear(method)
+	}
 	target.unlock()
 }
 
+func (o *Object) EdgeCount(direction EdgeDirection) int {
+	o.lock()
+	result := len(o.edges[direction])
+	o.unlock()
+	return result
+}
+
+func (o *Object) EdgeIterator(direction EdgeDirection, ei func(target *Object, edge EdgeBitmap) bool) {
+	o.lock()
+	for target, edge := range o.edges[direction] {
+		o.unlock()
+		if !ei(target, edge) {
+			return
+		}
+		o.lock()
+	}
+	o.unlock()
+}
+
 func (o *Object) ChildOf(parent *Object) {
 	o.lock()
 	if o.parent != nil {

diff --git a/modules/integrations/activedirectory/analyze/analyze-ad.go b/modules/integrations/activedirectory/analyze/analyze-ad.go
@@ -678,12 +678,13 @@ func init() {
 				}
 
 				// Add the DCsync combination flag
-				for p, methods := range o.PwnableBy {
-					if methods.IsSet(activedirectory.EdgeDSReplicationGetChanges) && methods.IsSet(activedirectory.EdgeDSReplicationGetChangesAll) {
+				o.EdgeIterator(engine.In, func(target *engine.Object, edge engine.EdgeBitmap) bool {
+					if edge.IsSet(activedirectory.EdgeDSReplicationGetChanges) && edge.IsSet(activedirectory.EdgeDSReplicationGetChangesAll) {
 						// DCsync attack WOT WOT
-						p.EdgeTo(o, activedirectory.EdgeDCsync)
+						target.EdgeTo(o, activedirectory.EdgeDCsync)
 					}
-				}
+					return true
+				})
 			},
 		},
 	)
@@ -1218,12 +1219,13 @@ func init() {
 
 			// Find the computer AD object if any
 			var computer *engine.Object
-			for target, edges := range machine.CanPwn {
-				if edges.IsSet(EdgeAuthenticatesAs) && target.Type() == engine.ObjectTypeComputer {
+			machine.EdgeIterator(engine.Out, func(target *engine.Object, edge engine.EdgeBitmap) bool {
+				if edge.IsSet(EdgeAuthenticatesAs) && target.Type() == engine.ObjectTypeComputer {
 					computer = target
-					break
+					return false //break
 				}
-			}
+				return true
+			})
 
 			if computer == nil {
 				continue
@@ -1452,7 +1454,7 @@ func init() {
 				if sources, found := ao.FindMulti(engine.ObjectSid, engine.AttributeValueSID(sid)); found {
 					for _, source := range sources {
 						if source.Type() != engine.ObjectTypeForeignSecurityPrincipal {
-							source.PwnsEx(foreign, activedirectory.EdgeForeignIdentity, true)
+							source.EdgeToEx(foreign, activedirectory.EdgeForeignIdentity, true)
 						}
 					}
 				}
@@ -1469,20 +1471,19 @@ func init() {
 		for _, gpo := range ao.Filter(func(o *engine.Object) bool {
 			return o.Type() == engine.ObjectTypeGroupPolicyContainer
 		}).Slice() {
-			for group, methods := range gpo.PwnableBy {
-
+			gpo.EdgeIterator(engine.In, func(group *engine.Object, methods engine.EdgeBitmap) bool {
 				groupname := group.OneAttrString(engine.SAMAccountName)
 				if strings.Contains(groupname, "%") {
 					// Lowercase for ease
 					groupname := strings.ToLower(groupname)
 
 					// It has some sort of % variable in it, let's go
-					for affected, amethods := range gpo.CanPwn {
+					gpo.EdgeIterator(engine.Out, func(affected *engine.Object, amethods engine.EdgeBitmap) bool {
 						if amethods.IsSet(activedirectory.EdgeAffectedByGPO) && affected.Type() == engine.ObjectTypeComputer {
 							netbiosdomain, computername, found := strings.Cut(affected.OneAttrString(engine.DownLevelLogonName), "\\")
 							if !found {
 								ui.Error().Msgf("Could not parse downlevel logon name %v", affected.OneAttrString(engine.DownLevelLogonName))
-								continue
+								return true //continue
 							}
 							computername = strings.TrimRight(computername, "$")
 
@@ -1507,7 +1508,7 @@ func init() {
 								warnlines++
 							} else if len(targetgroups) == 1 {
 								for _, edge := range methods.Edges() {
-									targetgroups[0].PwnsEx(affected, edge, true)
+									targetgroups[0].EdgeToEx(affected, edge, true)
 								}
 							} else {
 								ui.Warn().Msgf("Found multiple groups for %v: %v", realgroup, targetgroups)
@@ -1516,10 +1517,11 @@ func init() {
 								}
 							}
 						}
-					}
+						return true
+					})
 				}
-
-			}
+				return true
+			})
 		}
 		if warnlines > 0 {
 			ui.Warn().Msgf("%v groups could not be resolved, this could affect analysis results", warnlines)

diff --git a/modules/integrations/localmachine/analyze/postprocessing.go b/modules/integrations/localmachine/analyze/postprocessing.go
@@ -13,7 +13,7 @@ func init() {
 		for _, o := range ao.Slice() {
 			if o.HasAttrValue(engine.MetaDataSource, ln) {
 				if o.HasAttr(activedirectory.ObjectSid) {
-					if len(o.CanPwn) == 0 && len(o.PwnableBy) == 0 {
+					if o.EdgeCount(engine.Out)+o.EdgeCount(engine.In) == 0 {
 						ui.Debug().Msgf("Object has no graph connections: %v", o.Label())
 					}
 					warns++