Cleanup Code

Add debug option to load options Run which better
lewiswu1209 · Nov 4, 2019 · 1ccdc5f · 1ccdc5f
1 parent 4b54121
commit 1ccdc5f
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 30 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -7,7 +7,7 @@ plugins {
     kotlin("jvm") version "1.3.50"
 }
 
-group = "com.wavProductions.www"
+group = "com.wavproductions.www"
 version = "0.1-SNAPSHOT"
 
 repositories {

diff --git a/src/main/kotlin/com/wavproductions/www/armitage/Main.kt b/src/main/kotlin/com/wavproductions/www/armitage/Main.kt
@@ -7,16 +7,17 @@ import java.io.InputStream
 import java.nio.file.Files
 import java.nio.file.Paths
 
-fun main(args: MutableList<String>) {
+fun main(args: Array<String>) {
     val resources = Paths.get("resources").toAbsolutePath()
     if (Files.notExists(resources)) {
         Files.createDirectories(resources)
     }
     loadConfig()
     val conn = Metasploit()
-    val debug = args.contains("debug") //testing metasploit connection and calls
+    val debug = args.contains("debug-cmd") //testing metasploit connection and calls
     if (debug) {
-        conn.load()
+        conn.load(debug = true)
+
     } else {
         val iconPath = Paths.get(resources.toString(), "icon.png")
         if (Files.notExists(iconPath)) {

diff --git a/src/main/kotlin/com/wavproductions/www/armitage/Metasploit.kt b/src/main/kotlin/com/wavproductions/www/armitage/Metasploit.kt
@@ -2,22 +2,16 @@ package com.wavproductions.www.armitage
 
 import java.io.Closeable
 import java.io.IOException
-import java.io.InputStream
 import java.net.InetAddress
 import java.net.Socket
-import java.nio.file.Files
-import java.nio.file.Path
 import java.security.InvalidParameterException
 import javax.net.ssl.SSLSocketFactory
 
 class Metasploit : Closeable {
     private var console: Process?
-    private val temp: Path = Files.createTempFile("console", "MCArmitage")
-    private val fileIn: InputStream
     private var rpc: Socket? = null //will be a direct connection, so most likely a socket
 
     init {
-        fileIn = Files.newInputStream(temp)
         console = null
     }
 
@@ -28,7 +22,8 @@ class Metasploit : Closeable {
         port: Int = 55552,
         ssl: Boolean = false,
         program: String? = null,
-        local: Boolean = true
+        local: Boolean = true,
+        debug: Boolean = false
     ) {
         if (port < 0 || port > 65535) {
             throw InvalidParameterException("port is invalid! Valid values are 0-65535")
@@ -37,15 +32,25 @@ class Metasploit : Closeable {
         try {
             connected = attemptConnect(username, password, ip, port, ssl)
         } catch (ignored: Exception) {
+            if (debug) {
+                ignored.printStackTrace()
+            }
         }
         if (!connected && local) {
             val location = program ?: locateMetasploit() ?: throw NullPointerException("Unable to locate console!")
             val builder = ProcessBuilder(location)
             builder.redirectOutput(ProcessBuilder.Redirect.PIPE)
             builder.redirectError(ProcessBuilder.Redirect.PIPE)
+            if (debug) {
+                builder.redirectOutput(ProcessBuilder.Redirect.INHERIT)
+                builder.redirectError(ProcessBuilder.Redirect.INHERIT)
+            }
             builder.redirectInput(ProcessBuilder.Redirect.PIPE)
             console = builder.start()
-            //start rpc
+            //await startup
+            console?.outputStream?.write("load msgrpc ServerHost='${ip.hostAddress}' ServerPort=$port User='$username' Pass='$password'\n".toByteArray())
+            console?.outputStream?.flush()
+            //await rpc start
             connected = attemptConnect(username, password, ip, port, ssl)
         }
         if (!connected) {
@@ -55,11 +60,25 @@ class Metasploit : Closeable {
 
     private fun locateMetasploit(): String? {
         val builder = ProcessBuilder("which", "msfconsole") //linux lookup for now... might add windows later
-        builder.redirectOutput(ProcessBuilder.Redirect.INHERIT)
-        builder.redirectError(ProcessBuilder.Redirect.INHERIT)
-        builder.redirectInput(ProcessBuilder.Redirect.DISCARD)
-        console = builder.start()
-        return null
+        builder.redirectOutput(ProcessBuilder.Redirect.PIPE)
+        builder.redirectError(ProcessBuilder.Redirect.PIPE)
+        builder.redirectInput(ProcessBuilder.Redirect.PIPE)
+        val which = builder.start()
+        while (which.isAlive) {
+            Thread.onSpinWait()
+        }
+        if (String(which.errorStream.readAllBytes()).contains("which: no")) {
+            return null
+        }
+        val programs = String(which.inputStream.readAllBytes()).split("\n")
+        var selected: String? = null
+        for (x in programs) {
+            if (x.endsWith("msfconsole")) { //select first program located by which
+                selected = x
+                break
+            }
+        }
+        return selected
     }
 
     private fun attemptConnect(username: String, password: String, ip: InetAddress, port: Int, ssl: Boolean): Boolean {
@@ -82,17 +101,6 @@ class Metasploit : Closeable {
         return false //rpc failed to init!
     }
 
-    fun sendConsoleCommand(command: ByteArray, flush: Boolean = false) {
-        console?.outputStream?.write(command)
-        if (flush) {
-            console?.outputStream?.flush()
-        }
-    }
-
-    fun readConsole(): ByteArray {
-        return fileIn.readAllBytes()
-    }
-
     fun readRPC(): ByteArray {
         return rpc?.inputStream?.readAllBytes() ?: ByteArray(0)
     }
@@ -107,7 +115,5 @@ class Metasploit : Closeable {
     override fun close() { //destroy resources
         console?.destroyForcibly()
         rpc?.close()
-        fileIn.close()
-        Files.deleteIfExists(temp)
     }
 }