4141import burp .api .montoya .proxy .websocket .ProxyWebSocketCreationHandler ;
4242import burp .api .montoya .ui .editor .extension .ExtensionProvidedWebSocketMessageEditor ;
4343import burp .api .montoya .ui .editor .extension .WebSocketMessageEditorProvider ;
44+ import burp .api .montoya .websocket .BinaryMessage ;
45+ import burp .api .montoya .websocket .BinaryMessageAction ;
46+ import burp .api .montoya .websocket .Direction ;
47+ import burp .api .montoya .websocket .MessageHandler ;
48+ import burp .api .montoya .websocket .TextMessage ;
49+ import burp .api .montoya .websocket .TextMessageAction ;
50+ import burp .api .montoya .websocket .WebSocketCreated ;
51+ import burp .api .montoya .websocket .WebSocketCreatedHandler ;
4452import java .awt .Component ;
4553import java .awt .event .MouseAdapter ;
4654import java .awt .event .MouseEvent ;
@@ -1825,27 +1833,65 @@ private HttpMessage replaceProxyMessage(
18251833
18261834 }
18271835
1828- protected final class WebSocketCreationHander implements ProxyWebSocketCreationHandler {
1836+ protected final class WebSocketCreationHander implements ProxyWebSocketCreationHandler , WebSocketCreatedHandler {
18291837
18301838 private final MontoyaApi api ;
18311839
18321840 public WebSocketCreationHander (MontoyaApi api ) {
18331841 this .api = api ;
18341842 api .proxy ().registerWebSocketCreationHandler (this );
1843+ api .websockets ().registerWebSocketCreatedHandler (this );
18351844 }
18361845
18371846 @ Override
18381847 public void handleWebSocketCreation (ProxyWebSocketCreation proxyWebSocketCreation ) {
1839- proxyWebSocketCreation .proxyWebSocket ().registerProxyMessageHandler (new WebSocktHander (api , proxyWebSocketCreation ));
1848+ proxyWebSocketCreation .proxyWebSocket ().registerProxyMessageHandler (new WebSocktProxyMessageHander (api , proxyWebSocketCreation ));
1849+ }
1850+
1851+ @ Override
1852+ public void handleWebSocketCreated (WebSocketCreated webSocketCreated ) {
1853+ webSocketCreated .webSocket ().registerMessageHandler (new WebSocktMessageHander (api , webSocketCreated ));
18401854 }
18411855
18421856 }
18431857
1844- protected final class WebSocktHander implements ProxyMessageHandler {
1858+ protected final class WebSocktMessageHander implements MessageHandler {
1859+ private final MontoyaApi api ;
1860+ private final WebSocketCreated webSocketCreated ;
1861+
1862+ public WebSocktMessageHander (MontoyaApi api , WebSocketCreated webSocketCreated ) {
1863+ this .api = api ;
1864+ this .webSocketCreated = webSocketCreated ;
1865+ }
1866+
1867+ @ Override
1868+ public TextMessageAction handleTextMessage (TextMessage textMessage ) {
1869+ // WebSockt 出力
1870+ if (getProperty ().getLoggingProperty ().isAutoLogging () && getProperty ().getLoggingProperty ().isWebSocketLog ()) {
1871+ ToolSource toolSource = webSocketCreated .toolSource ();
1872+ logging .writeWebSocketToolMessage (toolSource .toolType (), webSocketCreated , textMessage );
1873+ }
1874+ return TextMessageAction .continueWith (textMessage );
1875+ }
1876+
1877+ @ Override
1878+ public BinaryMessageAction handleBinaryMessage (BinaryMessage binaryMessage ) {
1879+ // WebSockt 出力
1880+ if (getProperty ().getLoggingProperty ().isAutoLogging () && getProperty ().getLoggingProperty ().isWebSocketLog ()) {
1881+ ToolSource toolSource = webSocketCreated .toolSource ();
1882+ logging .writeWebSocektToolMessage (toolSource .toolType (), webSocketCreated , binaryMessage );
1883+ }
1884+ return BinaryMessageAction .continueWith (binaryMessage );
1885+ }
1886+
1887+ }
1888+
1889+
1890+ protected final class WebSocktProxyMessageHander implements ProxyMessageHandler {
18451891 private final MontoyaApi api ;
18461892 private final ProxyWebSocketCreation proxyWebSocketCreation ;
18471893
1848- public WebSocktHander (MontoyaApi api , ProxyWebSocketCreation proxyWebSocketCreation ) {
1894+ public WebSocktProxyMessageHander (MontoyaApi api , ProxyWebSocketCreation proxyWebSocketCreation ) {
18491895 this .api = api ;
18501896 this .proxyWebSocketCreation = proxyWebSocketCreation ;
18511897 proxyWebSocketCreation .proxyWebSocket ().registerProxyMessageHandler (this );
@@ -1854,35 +1900,39 @@ public WebSocktHander(MontoyaApi api, ProxyWebSocketCreation proxyWebSocketCreat
18541900 @ Override
18551901 public TextMessageReceivedAction handleTextMessageReceived (InterceptedTextMessage interceptedTextMessage ) {
18561902 // WebSockt 出力
1857- if (getProperty ().getLoggingProperty ().isAutoLogging () && getProperty ().getLoggingProperty ().isWebSocketLog ()) {
1858- logging .writeWebSocktFinalMessage (proxyWebSocketCreation , interceptedTextMessage );
1903+ if (getProperty ().getLoggingProperty ().isAutoLogging () && getProperty ().getLoggingProperty ().isWebSocketLog () &&
1904+ interceptedTextMessage .direction () == Direction .SERVER_TO_CLIENT ) {
1905+ logging .writeWebSocketFinalMessage (this .proxyWebSocketCreation , interceptedTextMessage );
18591906 }
18601907 return TextMessageReceivedAction .continueWith (interceptedTextMessage );
18611908 }
18621909
18631910 @ Override
18641911 public TextMessageToBeSentAction handleTextMessageToBeSent (InterceptedTextMessage interceptedTextMessage ) {
18651912 // WebSockt 出力
1866- if (getProperty ().getLoggingProperty ().isAutoLogging () && getProperty ().getLoggingProperty ().isWebSocketLog ()) {
1867- logging .writeWebSocktFinalMessage (proxyWebSocketCreation , interceptedTextMessage );
1913+ if (getProperty ().getLoggingProperty ().isAutoLogging () && getProperty ().getLoggingProperty ().isWebSocketLog () &&
1914+ interceptedTextMessage .direction () == Direction .CLIENT_TO_SERVER ) {
1915+ logging .writeWebSocketFinalMessage (proxyWebSocketCreation , interceptedTextMessage );
18681916 }
18691917 return TextMessageToBeSentAction .continueWith (interceptedTextMessage );
18701918 }
18711919
18721920 @ Override
18731921 public BinaryMessageReceivedAction handleBinaryMessageReceived (InterceptedBinaryMessage interceptedBinaryMessage ) {
18741922 // WebSockt 出力
1875- if (getProperty ().getLoggingProperty ().isAutoLogging () && getProperty ().getLoggingProperty ().isWebSocketLog ()) {
1876- logging .writeWebSocktFinalMessage (proxyWebSocketCreation , interceptedBinaryMessage );
1923+ if (getProperty ().getLoggingProperty ().isAutoLogging () && getProperty ().getLoggingProperty ().isWebSocketLog () &&
1924+ interceptedBinaryMessage .direction () == Direction .SERVER_TO_CLIENT ) {
1925+ logging .writeWebSocketFinalMessage (proxyWebSocketCreation , interceptedBinaryMessage );
18771926 }
18781927 return BinaryMessageReceivedAction .continueWith (interceptedBinaryMessage );
18791928 }
18801929
18811930 @ Override
18821931 public BinaryMessageToBeSentAction handleBinaryMessageToBeSent (InterceptedBinaryMessage interceptedBinaryMessage ) {
18831932 // WebSockt 出力
1884- if (getProperty ().getLoggingProperty ().isAutoLogging () && getProperty ().getLoggingProperty ().isWebSocketLog ()) {
1885- logging .writeWebSocktFinalMessage (proxyWebSocketCreation , interceptedBinaryMessage );
1933+ if (getProperty ().getLoggingProperty ().isAutoLogging () && getProperty ().getLoggingProperty ().isWebSocketLog () &&
1934+ interceptedBinaryMessage .direction () == Direction .CLIENT_TO_SERVER ) {
1935+ logging .writeWebSocketFinalMessage (proxyWebSocketCreation , interceptedBinaryMessage );
18861936 }
18871937 return BinaryMessageToBeSentAction .continueWith (interceptedBinaryMessage );
18881938 }
0 commit comments