burp-blazor/src/main/kotlin/BurpExtender.kt

97 lines
3.0 KiB
Kotlin

package burp
import burp.api.montoya.BurpExtension
import burp.api.montoya.MontoyaApi
import burp.api.montoya.core.ByteArray
import burp.api.montoya.websocket.*
const val EXTENSION_NAME = "Burp Blazor"
val openSockets = mutableMapOf<Int, WebSocket>()
class BurpExtender : BurpExtension {
private lateinit var api: MontoyaApi
private lateinit var ui: UI
override fun initialize(api: MontoyaApi) {
log.registerStreams(api.logging().output(), api.logging().error())
log.setLevel(Logger.Level.DEBUG)
this.api = api
api.extension().setName(EXTENSION_NAME)
ui = UI(api)
val webSockConnHandler = WSCH()
webSockConnHandler.ui = ui
api.websockets().registerWebSocketCreatedHandler(webSockConnHandler)
log.info("Initialized $EXTENSION_NAME")
}
}
class WSH : MessageHandler {
var handlerID = 0
private var handshakeDone = false
lateinit var ui: UI
lateinit var socket: WebSocket
override fun handleTextMessage(textMessage: TextMessage): TextMessageAction {
if (!handshakeDone && textMessage.direction() == Direction.CLIENT_TO_SERVER) {
if ("\"protocol\":\"blazorpack\"" in textMessage.payload()) {
log.debug("[$handlerID] HANDSHAKE: $textMessage")
openSockets[handlerID] = socket
handshakeDone = true
}
}
return TextMessageAction.continueWith(textMessage)
}
override fun handleBinaryMessage(binaryMessage: BinaryMessage): BinaryMessageAction {
if (handshakeDone) {
try {
if (binaryMessage.direction() == Direction.CLIENT_TO_SERVER) {
return handleOutBinaryMessage(binaryMessage.payload())
}
return handleInBinaryMessage(binaryMessage.payload())
} catch (e: Exception) {
log.debug(e.stackTraceToString())
}
}
return BinaryMessageAction.continueWith(binaryMessage)
}
private fun handleInBinaryMessage(binaryMessage: ByteArray): BinaryMessageAction {
ui.appendMessages(handlerID, "<- To client", decode(binaryMessage.bytes))
return BinaryMessageAction.continueWith(binaryMessage)
}
private fun handleOutBinaryMessage(binaryMessage: ByteArray): BinaryMessageAction {
ui.appendMessages(handlerID, "-> To server", decode(binaryMessage.bytes))
return BinaryMessageAction.continueWith(binaryMessage)
}
override fun onClose() {
if (handshakeDone) {
log.debug("Websocket with ID $handlerID closed")
openSockets.remove(handlerID)
}
super.onClose()
}
}
class WSCH : WebSocketCreatedHandler {
private var handlerCount = 0
lateinit var ui: UI
override fun handleWebSocketCreated(webSocketCreated: WebSocketCreated) {
val handler = WSH()
handler.handlerID = handlerCount
handlerCount++
handler.ui = ui
handler.socket = webSocketCreated.webSocket()
handler.socket.registerMessageHandler(handler)
}
}