Implemented memory scan module dump and log saving to file

ES.Shed/AnalysisReport.fs

@@ -10,10 +10,11 @@ open System.Security.Cryptography
 open Microsoft.Diagnostics.Runtime
 open System.Runtime.Serialization.Json
 
-type AnalysisReport(messageBus: MessageBus) =
+type AnalysisReport(settings: HandlerSettings) =
     let _messages = new List<IMessage>()
-    let (trace, info, _, error) = createLoggers(messageBus)
-
+    let (trace, info, _, error) = createLoggers(settings.MessageBus)
+    let _logStringBuilder = new StringBuilder()
+
     let md5(buffer: Byte array) =
         use md5 = MD5.Create()
         let hash = md5.ComputeHash(buffer)
@@ -35,7 +36,6 @@ type AnalysisReport(messageBus: MessageBus) =
         let mutable name = inName
         if String.IsNullOrWhiteSpace(name) || not <| Uri.IsWellFormedUriString(name, UriKind.RelativeOrAbsolute) then
             name <- md5(Encoding.Default.GetBytes(name))
-
             if isDll then name <- name + ".dll"
             elif isExe then name <- name + ".exe"
 
@@ -46,13 +46,29 @@ type AnalysisReport(messageBus: MessageBus) =
         info("Saved dynamic module: " + name)
 
     let saveModule(outputDir: String, modEvent: ExtractedManagedModuleEvent) =
-        let name =
-            if String.IsNullOrWhiteSpace(modEvent.Module.Name) then Guid.NewGuid().ToString("N")
-            else
-                let chunks = modEvent.Module.Name.Split(',')
-                if chunks.Length > 0 then Path.GetFileName(chunks.[0])
-                else modEvent.Module.Name
-        saveFileBuffer(outputDir, name, modEvent.Bytes, modEvent.IsDll, modEvent.isExecutable)
+        let mutable name =
+            match modEvent.Assembly with
+            | Some assembly when assembly.ManifestModule <> null -> assembly.ManifestModule.ScopeName
+            | _ -> 
+                let tmpName = 
+                    if String.IsNullOrWhiteSpace(modEvent.Module.Name) then Guid.NewGuid().ToString("N")
+                    else
+                        let chunks = modEvent.Module.Name.Split(',')
+                        if chunks.Length > 0 then Path.GetFileName(chunks.[0])
+                        else modEvent.Module.Name
+                if String.IsNullOrWhiteSpace(Path.GetExtension(tmpName)) then
+                    if modEvent.IsDll then tmpName + ".dll"
+                    else tmpName + ".exe"
+                else tmpName
+
+        saveFileBuffer(outputDir, name, modEvent.Bytes, modEvent.IsDll, modEvent.IsExecutable)
+
+    let saveMemoryScanModule(outputDir: String, modEvent: ExtractedManagedModuleViaMemoryScanEvent) =
+        let name = 
+            match modEvent.Assembly with
+            | Some assembly when assembly.ManifestModule <> null -> assembly.ManifestModule.ScopeName
+            | _ -> Guid.NewGuid().ToString("N") + if modEvent.IsDll then ".dll" else ".exe"
+        saveFileBuffer(outputDir, name, modEvent.Bytes, modEvent.IsDll, modEvent.IsExecutable)
 
     let saveProcessModule(outputDir: String, modEvent: ExtractedProcessModule) =
         let processModule = modEvent.Module
@@ -89,28 +105,41 @@ type AnalysisReport(messageBus: MessageBus) =
         Thread.CurrentThread.CurrentCulture <- currentCulture
         info("Heap json content saved to file: heap.json")
 
+    let handleLogMessageEvent(event: LogMessageEvent) = 
+        let msg = String.Format("[{0}] {1}", event.Level, event.Message)       
+        _logStringBuilder.AppendLine(msg) |> ignore
+
     let generateReport(command: GenerateReportCommand) =
         let outputDir = Path.Combine(Directory.GetCurrentDirectory(), "Result", command.ProcessId.ToString())
         Directory.CreateDirectory(outputDir) |> ignore
 
+        // handle all other messages
         for message in _messages do
             match message with
             | :? ExtractedManagedModuleEvent as modEvent -> saveModule(outputDir, modEvent)
             | :? HeapWalked as heapEvent -> serializeHeap(outputDir, heapEvent.Root)
             | :? ExtractedProcessModule as procModuleEvent -> saveProcessModule(outputDir, procModuleEvent)
+            | :? ExtractedManagedModuleViaMemoryScanEvent as memScanModEvent -> saveMemoryScanModule(outputDir, memScanModEvent)
             | _ -> ()
 
-        messageBus.Dispatch(new LogMessageEvent("Result saved to " + outputDir, LogLevel.Info))
+        info("Result saved to " + outputDir)
 
+        // save log to file
+        let logfile = Path.Combine(outputDir, "output.log")
+        File.WriteAllText(logfile, _logStringBuilder.ToString())
+
     member this.CanHandle(message: IMessage) =
         message :? ExtractedManagedModuleEvent ||
         message :? ExtractedProcessModule ||
         message :? GenerateReportCommand ||
+        message :? ExtractedManagedModuleViaMemoryScanEvent ||
+        message :? LogMessageEvent ||
         message :? HeapWalked
 
     member this.Handle(msg: IMessage) =
         match msg with
         | :? GenerateReportCommand as command -> generateReport(command)
+        | :? LogMessageEvent as moduleEvent -> handleLogMessageEvent(moduleEvent)
         | _ -> _messages.Add(msg)
 
     interface IMessageHandler with

ES.Shed/ES.Shed.fsproj

@@ -59,10 +59,12 @@
   <ItemGroup>
     <Compile Include="AssemblyInfo.fs" />
     <Content Include="packages.config" />
+    <Compile Include="Native.fs" />
     <Compile Include="Entities.fs" />
     <Compile Include="Messages.fs" />
     <Compile Include="IMessageHandler.fs" />
     <Compile Include="MessageBus.fs" />
+    <Compile Include="HandlerSettings.fs" />
     <Compile Include="MessageLogger.fs" />
     <Compile Include="ModuleDumper.fs" />
     <Compile Include="HeapDumper.fs" />

ES.Shed/HandlerSettings.fs

@@ -0,0 +1,6 @@
+namespace ES.Shed
+
+open System
+
+type HandlerSettings(messageBus: MessageBus) =
+    member val MessageBus = messageBus with get

ES.Shed/HeapDumper.fs

@@ -8,11 +8,11 @@ open System.Reflection
 open System.Collections.Generic
 open Microsoft.Diagnostics.Runtime
 
-type HeapDumper(messageBus: MessageBus) =
+type HeapDumper(settings: HandlerSettings) =
     let _objectsAlreadyAnalyzed = new HashSet<UInt64>()
     let _basePath = Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location)
     let _loggedStrings = new HashSet<String>()
-    let (trace, info, _, error) = createLoggers(messageBus)
+    let (trace, info, _, error) = createLoggers(settings.MessageBus)
 
     let isInterestingString(o: HeapObject) =
         if o.Value <> null && o.Value.ToString().Length >= 10 then
@@ -185,7 +185,7 @@ type HeapDumper(messageBus: MessageBus) =
             for objAddr in heap.EnumerateObjectAddresses() do
                 analyzeObjectAddress(objAddr, heap, root, None)
 
-            messageBus.Dispatch(new HeapWalked(root))
+            settings.MessageBus.Dispatch(new HeapWalked(root))
             info("Heap dump completed")
         else
             error("Heap is not walkable")

ES.Shed/MessageLogger.fs

@@ -1,6 +1,7 @@
 namespace ES.Shed
 
 open System
+open System.Text
 open System.Collections.Generic
 open Microsoft.Diagnostics.Runtime
 
@@ -23,26 +24,29 @@ module LogginHelpers =
 
         (trace, info, warning, error)
 
-type MessageLogger(messageBus: MessageBus) =
+type MessageLogger(settings: HandlerSettings) =
     let mutable _logLevel = LogLevel.Info
 
+    let printWithColor(msg: String, color: ConsoleColor) =
+        Console.ForegroundColor <- color
+        Console.WriteLine(msg)
+        Console.ResetColor() 
+
     let trace(msg: String) =
-        Console.ForegroundColor <- ConsoleColor.DarkCyan
-        Console.WriteLine("[~] {0}", msg)
-        Console.ResetColor()        
-
-    let info(msg: String) =        
-        Console.WriteLine("[+] {0}", msg)
-
-    let warning(msg: String) =      
-        Console.ForegroundColor <- ConsoleColor.DarkYellow
-        Console.WriteLine("[-] {0}", msg)
-        Console.ResetColor()
-
-    let error(msg: String) =        
-        Console.ForegroundColor <- ConsoleColor.DarkRed
-        Console.WriteLine("[!] {0}", msg)
-        Console.ResetColor()
+        let m = "[~] " + msg
+        printWithColor(m, ConsoleColor.DarkCyan)
+
+    let info(msg: String) =
+        let m = "[+] " + msg
+        Console.WriteLine(msg)
+
+    let warning(msg: String) =
+        let m = "[-] " + msg
+        printWithColor(m, ConsoleColor.DarkYellow)
+
+    let error(msg: String) =
+        let m = "[!] " + msg
+        printWithColor(m, ConsoleColor.DarkRed)
 
     let _handlers = 
         [

ES.Shed/Messages.fs

@@ -1,6 +1,7 @@
 namespace ES.Shed
 
 open System
+open System.Reflection
 open Microsoft.Diagnostics
 open System.Diagnostics
 open Microsoft.Diagnostics.Runtime
@@ -46,12 +47,23 @@ type ExtractedManagedModuleEvent(clrModule: ClrModule, bytes: Byte array, isDll:
     member val Module = clrModule with get
     member val Bytes = bytes with get
     member val IsDll = isDll with get
-    member val isExecutable = isExec with get
+    member val IsExecutable = isExec with get
+    member val Assembly: Assembly option = None with get, set
 
     override this.ToString() =
         if not(String.IsNullOrWhiteSpace(clrModule.Name)) then clrModule.Name
         else String.Format("Unnamed module, len: {0}", this.Bytes.Length)
 
+type ExtractedManagedModuleViaMemoryScanEvent(bytes: Byte array, isDll: Boolean, isExec: Boolean) =
+    inherit BaseEvent()
+    member val Bytes = bytes with get
+    member val IsDll = isDll with get
+    member val IsExecutable = isExec with get
+    member val Assembly: Assembly option = None with get, set
+
+    override this.ToString() =
+        String.Format("Extracted module via memory scan, len: {0}", this.Bytes.Length)
+
 type ExtractedProcessModule(procModule: ProcessModule) =
     inherit BaseEvent()
     member val Module = procModule with get, set