From 38b45380f4642a366c23df563d3757680fdbb5c0 Mon Sep 17 00:00:00 2001
From: Spencer McIntyre <zeroSteiner@gmail.com>
Date: Wed, 19 May 2021 16:09:17 -0400
Subject: [PATCH] Fix and process block edges within the GraphML

---
 data/shellcode/block_api.x64.graphml |  9 +--------
 data/shellcode/block_api.x86.graphml |  8 +-------
 lib/rex/parser/graphml.rb            |  5 +++++
 lib/rex/payloads/shuffle.rb          | 17 +++++++++++------
 4 files changed, 18 insertions(+), 21 deletions(-)

diff --git a/data/shellcode/block_api.x64.graphml b/data/shellcode/block_api.x64.graphml
index f7f15da89e..7721dd7217 100644
--- a/data/shellcode/block_api.x64.graphml
+++ b/data/shellcode/block_api.x64.graphml
@@ -642,23 +642,16 @@
     <edge source="block.0x1000" target="block.0x1017"/>
     <edge source="block.0x1017" target="block.0x1023"/>
     <edge source="block.0x1023" target="block.0x102b"/>
-    <edge source="block.0x1023" target="block.0x102d"/>
     <edge source="block.0x102b" target="block.0x102d"/>
     <edge source="block.0x102d" target="block.0x1036"/>
-    <edge source="block.0x102d" target="block.0x1023"/>
     <edge source="block.0x1036" target="block.0x104b"/>
-    <edge source="block.0x1036" target="block.0x10bd"/>
     <edge source="block.0x104b" target="block.0x1056"/>
-    <edge source="block.0x104b" target="block.0x10bd"/>
     <edge source="block.0x1056" target="block.0x1064"/>
     <edge source="block.0x1064" target="block.0x1066"/>
-    <edge source="block.0x1064" target="block.0x10bc"/>
     <edge source="block.0x1066" target="block.0x1073"/>
-    <edge source="block.0x1073" target="block.0x1073"/>
     <edge source="block.0x1073" target="block.0x1082"/>
-    <edge source="block.0x1082" target="block.0x1064"/>
     <edge source="block.0x1082" target="block.0x108c"/>
+    <edge source="block.0x108c" target="block.0x10bc"/>
     <edge source="block.0x10bc" target="block.0x10bd"/>
-    <edge source="block.0x10bd" target="block.0x1017"/>
   </graph>
 </graphml>
diff --git a/data/shellcode/block_api.x86.graphml b/data/shellcode/block_api.x86.graphml
index 4779c0c569..356016d5d0 100644
--- a/data/shellcode/block_api.x86.graphml
+++ b/data/shellcode/block_api.x86.graphml
@@ -595,21 +595,15 @@
     <edge source="block.0x1000" target="block.0x100f"/>
     <edge source="block.0x100f" target="block.0x1018"/>
     <edge source="block.0x1018" target="block.0x101f"/>
-    <edge source="block.0x1018" target="block.0x1021"/>
     <edge source="block.0x101f" target="block.0x1021"/>
-    <edge source="block.0x1021" target="block.0x1018"/>
     <edge source="block.0x1021" target="block.0x1029"/>
     <edge source="block.0x1029" target="block.0x103a"/>
-    <edge source="block.0x1029" target="block.0x1086"/>
     <edge source="block.0x103a" target="block.0x1045"/>
     <edge source="block.0x1045" target="block.0x1049"/>
-    <edge source="block.0x1045" target="block.0x1085"/>
     <edge source="block.0x1049" target="block.0x1051"/>
-    <edge source="block.0x1051" target="block.0x1051"/>
     <edge source="block.0x1051" target="block.0x105d"/>
-    <edge source="block.0x105d" target="block.0x1045"/>
     <edge source="block.0x105d" target="block.0x1065"/>
+    <edge source="block.0x1065" target="block.0x1085"/>
     <edge source="block.0x1085" target="block.0x1086"/>
-    <edge source="block.0x1086" target="block.0x100f"/>
   </graph>
 </graphml>
diff --git a/lib/rex/parser/graphml.rb b/lib/rex/parser/graphml.rb
index 6872b15c2e..dc640f900f 100644
--- a/lib/rex/parser/graphml.rb
+++ b/lib/rex/parser/graphml.rb
@@ -362,6 +362,7 @@ module Rex
           def initialize(id)
             @id = id
             @edges = []
+            @subgraph = nil
             super()
           end
 
@@ -390,6 +391,9 @@ module Rex
           # @!attribute edges
           #   @return [Array] An array of all edges for which this node is either the source or the target.
           attr_reader :edges
+          # @!attribute subgraph
+          #   @return [Graph,nil] A subgraph contained within this node.
+          attr_accessor :subgraph
         end
       end
 
@@ -461,6 +465,7 @@ module Rex
 
           when 'graph'
             element = Element::Graph.from_xml_attributes(attrs)
+            @stack[-1].subgraph = element if @stack[-1].is_a? Element::Node
             @graphml.graphs << element
 
           when 'graphml'
diff --git a/lib/rex/payloads/shuffle.rb b/lib/rex/payloads/shuffle.rb
index 480a88fa6e..ade9acfac7 100644
--- a/lib/rex/payloads/shuffle.rb
+++ b/lib/rex/payloads/shuffle.rb
@@ -23,7 +23,7 @@ module Rex
       # @param name [String] An optional symbol name to apply to the assembly source.
       def self.from_graphml_file(file_path, arch: nil, name: nil)
         graphml = Rex::Parser::GraphML.from_file(file_path)
-        blocks = graphml.graphs.select { |graph| graph.attributes['type'] == 'block' }.sort_by { |graph| graph.attributes['address'] }
+        blocks = create_path(graphml.nodes.select { |_id,node| node.attributes['type'] == 'block' }, graphml.graphs[0].edges)
         blocks.map! { |block| { node: block, instructions: process_block(block) } }
 
         label_prefix = Rex::Text.rand_text_alpha_lower(4)
@@ -76,22 +76,27 @@ module Rex
         # contains nodes representing each of its instructions.
         #
         def process_block(block)
+          subgraph = block.subgraph
+          instructions = subgraph.nodes.select { |_id, node| node.attributes['type'] == 'instruction' }
+          create_path(instructions, subgraph.edges)
+        end
+
+        def create_path(nodes, edges)
           path = []
-          instructions = block.nodes.select { |_id, node| node.attributes['type'] == 'instruction' }
 
           # the initial choices are any node without a predecessor (dependency)
-          targets = block.edges.map(&:target)
-          choices = instructions.values.select { |node| !targets.include? node.id }
+          targets = edges.map(&:target)
+          choices = nodes.values.select { |node| !targets.include? node.id }
           until choices.empty?
             selection = choices.sample
             choices.delete(selection)
             path << selection
 
             # check each node for which the selection is a dependency
-            successors = selection.target_edges.map { |edge| instructions[edge.target] }
+            successors = selection.target_edges.map { |edge| nodes[edge.target] }
             successors.each do |successor|
               next if path.include? successor
-              next if !successor.source_edges.map { |edge| path.include? instructions[edge.source] }.all?
+              next if !successor.source_edges.map { |edge| path.include? nodes[edge.source] }.all?
 
               choices << successor
             end