Fix ImplantTemplate becomes null on GruntTable

cobbr · Sep 7, 2020 · 1cb750a · 1cb750a
1 parent 93eddf4
commit 1cb750a
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 8 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fix profile bug when HttpGetResponse differs from HttpPostResponse
 - Fix TaskKill display bug
 - Fix token impersonation issues
+- Fix ImplantTemplate becomes null on GruntTable
 
 ## [v0.6] - 2020-08-04
 ### Added

diff --git a/Covenant/Components/Grunts/GruntCommandCard.razor b/Covenant/Components/Grunts/GruntCommandCard.razor
@@ -99,6 +99,7 @@
         if (this.GruntCommand.CommandOutput != null && this.GruntCommand.CommandOutput.Id == commandOutput.Id)
         {
             this.GruntCommand = await Service.GetGruntCommand(this.GruntCommandId);
+            Service.DisposeContext();
             this.GruntCommand.CommandOutput ??= await Service.GetCommandOutput(this.GruntCommand.CommandOutputId);
             Service.DisposeContext();
             await this.InvokeAsync(() => this.StateHasChanged());
@@ -110,6 +111,7 @@
         if (this.GruntCommandId == gruntCommand.Id)
         {
             this.GruntCommand = await Service.GetGruntCommand(this.GruntCommandId);
+            Service.DisposeContext();
             this.GruntCommand.CommandOutput ??= await Service.GetCommandOutput(this.GruntCommand.CommandOutputId);
             Service.DisposeContext();
             await this.InvokeAsync(() => this.StateHasChanged());

diff --git a/Covenant/Components/Grunts/GruntIndex.razor b/Covenant/Components/Grunts/GruntIndex.razor
@@ -37,8 +37,8 @@
 <TabbedTerminals @ref="TabbedTerminals" @bind-TerminalGrunts="TerminalGrunts" OnRemoveTerminal="OnRemoveTerminal"/>
 
 @code {
-        [Parameter]
-        public List<Grunt> Grunts { get; set; }
+    [Parameter]
+    public List<Grunt> Grunts { get; set; }
     private List<Grunt> UnfilteredGrunts { get; set; }
 
     private TabbedTerminals TabbedTerminals { get; set; }
@@ -82,21 +82,26 @@
         }
     }
 
-    private void OnEditGrunt(object sender, Grunt grunt)
+    private async void OnEditGrunt(object sender, Grunt grunt)
     {
+        if (grunt.ImplantTemplate == null)
+        {
+            grunt.ImplantTemplate = await Service.GetImplantTemplate(grunt.ImplantTemplateId);
+            Service.DisposeContext();
+        }
         bool inDisplayedGrunts = this.Grunts.Any(G => G.Id == grunt.Id);
         if (inDisplayedGrunts)
         {
             this.Grunts[this.Grunts.FindIndex(G => G.Id == grunt.Id)] = grunt;
-            this.InvokeAsync(() => this.StateHasChanged());
+            await this.InvokeAsync(() => this.StateHasChanged());
         }
         if (this.UnfilteredGrunts.Any(G => G.Id == grunt.Id))
         {
             this.UnfilteredGrunts[this.UnfilteredGrunts.FindIndex(G => G.Id == grunt.Id)] = grunt;
             if (!inDisplayedGrunts && (!this.Hidden || grunt.Status != GruntStatus.Hidden) && grunt.Status != GruntStatus.Uninitialized)
             {
                 this.FilterGrunts();
-                this.InvokeAsync(() => this.StateHasChanged());
+                await this.InvokeAsync(() => this.StateHasChanged());
             }
         }
     }

diff --git a/Covenant/Components/Grunts/GruntTable.razor b/Covenant/Components/Grunts/GruntTable.razor
@@ -41,7 +41,7 @@
         <SortableTableHeader TItem="Grunt" TItem2="string" GetSortableProperty="(grunt) => grunt.Note">
             <HeaderContent>Note</HeaderContent>
         </SortableTableHeader>
-        <SortableTableHeader TItem="Grunt" TItem2="string" GetSortableProperty="(grunt) => grunt.ImplantTemplate.Name">
+        <SortableTableHeader TItem="Grunt" TItem2="string" GetSortableProperty="(grunt) => grunt.ImplantTemplate?.Name">
             <HeaderContent>Template</HeaderContent>
         </SortableTableHeader>
     </TableHeader>
@@ -66,7 +66,7 @@
             <td @attributes="InputAttributes">@grunt.LastCheckIn</td>
             <td @attributes="InputAttributes">@grunt.Status</td>
             <td @attributes="InputAttributes">@grunt.Note</td>
-            <td @attributes="InputAttributes">@grunt.ImplantTemplate.Name</td>
+            <td @attributes="InputAttributes">@grunt.ImplantTemplate?.Name</td>
         </tr>
     </TableRow>
     <ButtonTray>@ButtonTray</ButtonTray>
@@ -103,6 +103,7 @@
     protected async override Task OnInitializedAsync()
     {
         this.Grunts = this.Grunts ?? (await Service.GetGrunts()).ToList();
+        Service.DisposeContext();
     }
 
     private string GetGruntStatusOpacity(GruntStatus s)
@@ -120,7 +121,7 @@
     {
         return SearchTerm == string.Empty ||
             grunt.Name.Contains(SearchTerm, StringComparison.CurrentCultureIgnoreCase) ||
-            grunt.ImplantTemplate.Name.Contains(SearchTerm, StringComparison.CurrentCultureIgnoreCase) ||
+            (grunt.ImplantTemplate != null && grunt.ImplantTemplate.Name.Contains(SearchTerm, StringComparison.CurrentCultureIgnoreCase)) ||
             grunt.Hostname.Contains(SearchTerm, StringComparison.CurrentCultureIgnoreCase) ||
             grunt.UserName.Contains(SearchTerm, StringComparison.CurrentCultureIgnoreCase) ||
             grunt.Status.ToString().Contains(SearchTerm, StringComparison.CurrentCultureIgnoreCase) ||