From 1177575f77bf9a977ea5ec36a66e5055bf3042a6 Mon Sep 17 00:00:00 2001
From: DarkPhoenix <phoenix@mail.ru>
Date: Fri, 19 Apr 2019 23:12:29 +0300
Subject: [PATCH] Make item remove context menu safer

---
 gui/builtinContextMenus/itemRemove.py | 54 +++++++++++++++++++--------
 gui/fitCommands/gui/cargo/remove.py   |  2 +-
 2 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/gui/builtinContextMenus/itemRemove.py b/gui/builtinContextMenus/itemRemove.py
index b70ce3679..cd5b92015 100644
--- a/gui/builtinContextMenus/itemRemove.py
+++ b/gui/builtinContextMenus/itemRemove.py
@@ -39,35 +39,57 @@ class RemoveItem(ContextMenu):
             self.mainFrame.command.Submit(cmd.GuiChangeLocalModuleChargesCommand(
                 fitID=fitID, modules=selection, chargeItemID=None))
         elif srcContext == "droneItem":
-            self.mainFrame.command.Submit(cmd.GuiRemoveLocalDroneCommand(
-                fitID=fitID, position=fit.drones.index(selection[0]), amount=math.inf))
+            drone = selection[0]
+            if drone in fit.drones:
+                position = fit.drones.index(drone)
+                self.mainFrame.command.Submit(cmd.GuiRemoveLocalDroneCommand(
+                    fitID=fitID, position=position, amount=math.inf))
         elif srcContext == "fighterItem":
-            self.mainFrame.command.Submit(cmd.GuiRemoveLocalFighterCommand(
-                fitID=fitID, position=fit.fighters.index(selection[0])))
+            fighter = selection[0]
+            if fighter in fit.fighters:
+                position = fit.fighters.index(fighter)
+                self.mainFrame.command.Submit(cmd.GuiRemoveLocalFighterCommand(
+                    fitID=fitID, position=position))
         elif srcContext == "implantItem":
-            self.mainFrame.command.Submit(cmd.GuiRemoveImplantCommand(
-                fitID=fitID, position=fit.implants.index(selection[0])))
+            implant = selection[0]
+            if implant in fit.implants:
+                position = fit.implants.index(implant)
+                self.mainFrame.command.Submit(cmd.GuiRemoveImplantCommand(
+                    fitID=fitID, position=position))
         elif srcContext == "boosterItem":
-            self.mainFrame.command.Submit(cmd.GuiRemoveBoosterCommand(
-                fitID=fitID, position=fit.boosters.index(selection[0])))
+            booster = selection[0]
+            if booster in fit.boosters:
+                position = fit.boosters.index(booster)
+                self.mainFrame.command.Submit(cmd.GuiRemoveBoosterCommand(
+                    fitID=fitID, position=position))
         elif srcContext == "cargoItem":
+            cargo = selection[0]
             self.mainFrame.command.Submit(cmd.GuiRemoveCargoCommand(
-                fitID=fitID, itemID=selection[0].itemID))
+                fitID=fitID, itemID=cargo.itemID))
         elif srcContext == "projectedFit":
+            projectedFit = selection[0]
             self.mainFrame.command.Submit(cmd.GuiRemoveProjectedFitCommand(
-                fitID=fitID, projectedFitID=selection[0].ID))
+                fitID=fitID, projectedFitID=projectedFit.ID))
         elif srcContext == "projectedModule":
-            self.mainFrame.command.Submit(cmd.GuiRemoveProjectedModuleCommand(
-                fitID=fitID, position=fit.projectedModules.index(selection[0])))
+            mod = selection[0]
+            if mod in fit.projectedModules:
+                position = fit.projectedModules.index(mod)
+                self.mainFrame.command.Submit(cmd.GuiRemoveProjectedModuleCommand(
+                    fitID=fitID, position=position))
         elif srcContext == "projectedDrone":
+            drone = selection[0]
             self.mainFrame.command.Submit(cmd.GuiRemoveProjectedDroneCommand(
-                fitID=fitID, itemID=selection[0].itemID, amount=math.inf))
+                fitID=fitID, itemID=drone.itemID, amount=math.inf))
         elif srcContext == "projectedFighter":
-            self.mainFrame.command.Submit(cmd.GuiRemoveProjectedFighterCommand(
-                fitID=fitID, position=fit.projectedFighters.index(selection[0])))
+            fighter = selection[0]
+            if fighter in fit.projectedFighters:
+                position = fit.projectedFighters.index(fighter)
+                self.mainFrame.command.Submit(cmd.GuiRemoveProjectedFighterCommand(
+                    fitID=fitID, position=position))
         elif srcContext == "commandFit":
+            commandFit = selection[0]
             self.mainFrame.command.Submit(cmd.GuiRemoveCommandFitCommand(
-                fitID=fitID, commandFitID=selection[0].ID))
+                fitID=fitID, commandFitID=commandFit.ID))
 
 
 RemoveItem.register()
diff --git a/gui/fitCommands/gui/cargo/remove.py b/gui/fitCommands/gui/cargo/remove.py
index 8db09bcb9..6eede4d59 100644
--- a/gui/fitCommands/gui/cargo/remove.py
+++ b/gui/fitCommands/gui/cargo/remove.py
@@ -17,7 +17,7 @@ class GuiRemoveCargoCommand(wx.Command):
         self.itemID = itemID
 
     def Do(self):
-        cmd =CalcRemoveCargoCommand(fitID=self.fitID, cargoInfo=CargoInfo(itemID=self.itemID, amount=math.inf))
+        cmd = CalcRemoveCargoCommand(fitID=self.fitID, cargoInfo=CargoInfo(itemID=self.itemID, amount=math.inf))
         success = self.internalHistory.submit(cmd)
         wx.PostEvent(gui.mainFrame.MainFrame.getInstance(), GE.FitChanged(fitID=self.fitID))
         return success