Implement download_as and export as pwntools

5 files changed, 249 insertions, 34 deletions

diff --git a/application_router.go b/application_router.go
index 30ec7c6..334e9f3 100644
--- a/application_router.go
+++ b/application_router.go
@@ -283,12 +283,36 @@ func CreateApplicationRouter(applicationContext *ApplicationContext,
 				badRequest(c, err)
 				return
 			}
-			var format QueryFormat
+			var format GetMessageFormat
 			if err := c.ShouldBindQuery(&format); err != nil {
 				badRequest(c, err)
 				return
 			}
-			success(c, applicationContext.ConnectionStreamsController.GetConnectionPayload(c, id, format))
+
+			if messages, found := applicationContext.ConnectionStreamsController.GetConnectionMessages(c, id, format); !found {
+				notFound(c, gin.H{"connection": id})
+			} else {
+				success(c, messages)
+			}
+		})
+
+		api.GET("/streams/:id/download", func(c *gin.Context) {
+			id, err := RowIDFromHex(c.Param("id"))
+			if err != nil {
+				badRequest(c, err)
+				return
+			}
+			var format DownloadMessageFormat
+			if err := c.ShouldBindQuery(&format); err != nil {
+				badRequest(c, err)
+				return
+			}
+
+			if blob, found := applicationContext.ConnectionStreamsController.DownloadConnectionMessages(c, id, format); !found {
+				notFound(c, gin.H{"connection": id})
+			} else {
+				c.String(http.StatusOK, blob)
+			}
 		})
 
 		api.GET("/services", func(c *gin.Context) {
diff --git a/connection_streams_controller.go b/connection_streams_controller.go
index 9d73b0e..98f2aca 100644
--- a/connection_streams_controller.go
+++ b/connection_streams_controller.go
@@ -3,14 +3,19 @@ package main
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"github.com/eciavatta/caronte/parsers"
 	log "github.com/sirupsen/logrus"
+	"strings"
 	"time"
 )
 
-const InitialPayloadsSize = 1024
-const DefaultQueryFormatLimit = 8024
-const InitialRegexSlicesCount = 8
+const (
+	initialPayloadsSize     = 1024
+	defaultQueryFormatLimit = 8024
+	initialRegexSlicesCount = 8
+	pwntoolsMaxServerBytes  = 20
+)
 
 type ConnectionStream struct {
 	ID               RowID                   `bson:"_id"`
@@ -26,7 +31,7 @@ type ConnectionStream struct {
 
 type PatternSlice [2]uint64
 
-type Payload struct {
+type Message struct {
 	FromClient             bool             `json:"from_client"`
 	Content                string           `json:"content"`
 	Metadata               parsers.Metadata `json:"metadata"`
@@ -42,12 +47,17 @@ type RegexSlice struct {
 	To   uint64 `json:"to"`
 }
 
-type QueryFormat struct {
+type GetMessageFormat struct {
 	Format string `form:"format"`
 	Skip   uint64 `form:"skip"`
 	Limit  uint64 `form:"limit"`
 }
 
+type DownloadMessageFormat struct {
+	Format string `form:"format"`
+	Type   string `form:"type"`
+}
+
 type ConnectionStreamsController struct {
 	storage Storage
 }
@@ -58,13 +68,18 @@ func NewConnectionStreamsController(storage Storage) ConnectionStreamsController
 	}
 }
 
-func (csc ConnectionStreamsController) GetConnectionPayload(c context.Context, connectionID RowID,
-	format QueryFormat) []*Payload {
-	payloads := make([]*Payload, 0, InitialPayloadsSize)
+func (csc ConnectionStreamsController) GetConnectionMessages(c context.Context, connectionID RowID,
+	format GetMessageFormat) ([]*Message, bool) {
+	connection := csc.getConnection(c, connectionID)
+	if connection.ID.IsZero() {
+		return nil, false
+	}
+
+	payloads := make([]*Message, 0, initialPayloadsSize)
 	var clientIndex, serverIndex, globalIndex uint64
 
 	if format.Limit <= 0 {
-		format.Limit = DefaultQueryFormatLimit
+		format.Limit = defaultQueryFormatLimit
 	}
 
 	var clientBlocksIndex, serverBlocksIndex int
@@ -79,8 +94,8 @@ func (csc ConnectionStreamsController) GetConnectionPayload(c context.Context, c
 		return serverBlocksIndex < len(serverStream.BlocksIndexes)
 	}
 
-	var payload *Payload
-	payloadsBuffer := make([]*Payload, 0, 16)
+	var payload *Message
+	payloadsBuffer := make([]*Message, 0, 16)
 	contentChunkBuffer := new(bytes.Buffer)
 	var lastContentSlice []byte
 	var sideChanged, lastClient, lastServer bool
@@ -97,7 +112,7 @@ func (csc ConnectionStreamsController) GetConnectionPayload(c context.Context, c
 			}
 			size := uint64(end - start)
 
-			payload = &Payload{
+			payload = &Message{
 				FromClient:      true,
 				Content:         DecodeBytes(clientStream.Payload[start:end], format.Format),
 				Index:           start,
@@ -121,7 +136,7 @@ func (csc ConnectionStreamsController) GetConnectionPayload(c context.Context, c
 			}
 			size := uint64(end - start)
 
-			payload = &Payload{
+			payload = &Message{
 				FromClient:      false,
 				Content:         DecodeBytes(serverStream.Payload[start:end], format.Format),
 				Index:           start,
@@ -178,11 +193,118 @@ func (csc ConnectionStreamsController) GetConnectionPayload(c context.Context, c
 		if globalIndex > format.Skip+format.Limit {
 			// problem: the last chunk is not parsed, but can be ok because it is not finished
 			updateMetadata()
-			return payloads
+			return payloads, true
 		}
 	}
 
-	return payloads
+	return payloads, true
+}
+
+func (csc ConnectionStreamsController) DownloadConnectionMessages(c context.Context, connectionID RowID,
+	format DownloadMessageFormat) (string, bool) {
+	connection := csc.getConnection(c, connectionID)
+	if connection.ID.IsZero() {
+		return "", false
+	}
+
+	var sb strings.Builder
+	includeClient, includeServer := format.Type != "only_server", format.Type != "only_client"
+	isPwntools := format.Type == "pwntools"
+
+	var clientBlocksIndex, serverBlocksIndex int
+	var clientDocumentIndex, serverDocumentIndex int
+	var clientStream ConnectionStream
+	if includeClient {
+		clientStream = csc.getConnectionStream(c, connectionID, true, clientDocumentIndex)
+	}
+	var serverStream ConnectionStream
+	if includeServer {
+		serverStream = csc.getConnectionStream(c, connectionID, false, serverDocumentIndex)
+	}
+
+	hasClientBlocks := func() bool {
+		return clientBlocksIndex < len(clientStream.BlocksIndexes)
+	}
+	hasServerBlocks := func() bool {
+		return serverBlocksIndex < len(serverStream.BlocksIndexes)
+	}
+
+	if isPwntools {
+		if format.Format == "base32" || format.Format == "base64" {
+			sb.WriteString("import base64\n")
+		}
+		sb.WriteString("from pwn import *\n\n")
+		sb.WriteString(fmt.Sprintf("p = remote('%s', %d)\n", connection.DestinationIP, connection.DestinationPort))
+	}
+
+	lastIsClient, lastIsServer := true, true
+	for !clientStream.ID.IsZero() || !serverStream.ID.IsZero() {
+		if hasClientBlocks() && (!hasServerBlocks() || // next payload is from client
+			clientStream.BlocksTimestamps[clientBlocksIndex].UnixNano() <=
+				serverStream.BlocksTimestamps[serverBlocksIndex].UnixNano()) {
+			start := clientStream.BlocksIndexes[clientBlocksIndex]
+			end := 0
+			if clientBlocksIndex < len(clientStream.BlocksIndexes)-1 {
+				end = clientStream.BlocksIndexes[clientBlocksIndex+1]
+			} else {
+				end = len(clientStream.Payload)
+			}
+
+			if !lastIsClient {
+				sb.WriteString("\n")
+			}
+			lastIsClient = true
+			lastIsServer = false
+			if isPwntools {
+				sb.WriteString(decodePwntools(clientStream.Payload[start:end], true, format.Format))
+			} else {
+				sb.WriteString(DecodeBytes(clientStream.Payload[start:end], format.Format))
+			}
+			clientBlocksIndex++
+		} else { // next payload is from server
+			start := serverStream.BlocksIndexes[serverBlocksIndex]
+			end := 0
+			if serverBlocksIndex < len(serverStream.BlocksIndexes)-1 {
+				end = serverStream.BlocksIndexes[serverBlocksIndex+1]
+			} else {
+				end = len(serverStream.Payload)
+			}
+
+			if !lastIsServer {
+				sb.WriteString("\n")
+			}
+			lastIsClient = false
+			lastIsServer = true
+			if isPwntools {
+				sb.WriteString(decodePwntools(serverStream.Payload[start:end], false, format.Format))
+			} else {
+				sb.WriteString(DecodeBytes(serverStream.Payload[start:end], format.Format))
+			}
+			serverBlocksIndex++
+		}
+
+		if includeClient && !hasClientBlocks() {
+			clientDocumentIndex++
+			clientBlocksIndex = 0
+			clientStream = csc.getConnectionStream(c, connectionID, true, clientDocumentIndex)
+		}
+		if includeServer && !hasServerBlocks() {
+			serverDocumentIndex++
+			serverBlocksIndex = 0
+			serverStream = csc.getConnectionStream(c, connectionID, false, serverDocumentIndex)
+		}
+	}
+
+	return sb.String(), true
+}
+
+func (csc ConnectionStreamsController) getConnection(c context.Context, connectionID RowID) Connection {
+	var connection Connection
+	if err := csc.storage.Find(Connections).Context(c).Filter(OrderedDocument{{"_id", connectionID}}).
+		First(&connection); err != nil {
+		log.WithError(err).WithField("id", connectionID).Panic("failed to get connection")
+	}
+	return connection
 }
 
 func (csc ConnectionStreamsController) getConnectionStream(c context.Context, connectionID RowID, fromClient bool,
@@ -199,7 +321,7 @@ func (csc ConnectionStreamsController) getConnectionStream(c context.Context, co
 }
 
 func findMatchesBetween(patternMatches map[uint][]PatternSlice, from, to uint64) []RegexSlice {
-	regexSlices := make([]RegexSlice, 0, InitialRegexSlicesCount)
+	regexSlices := make([]RegexSlice, 0, initialRegexSlicesCount)
 	for _, slices := range patternMatches {
 		for _, slice := range slices {
 			if from > slice[1] || to <= slice[0] {
@@ -225,3 +347,27 @@ func findMatchesBetween(patternMatches map[uint][]PatternSlice, from, to uint64)
 	}
 	return regexSlices
 }
+
+func decodePwntools(payload []byte, isClient bool, format string) string {
+	if !isClient && len(payload) > pwntoolsMaxServerBytes {
+		payload = payload[len(payload)-pwntoolsMaxServerBytes:]
+	}
+
+	var content string
+	switch format {
+	case "hex":
+		content = fmt.Sprintf("bytes.fromhex('%s')", DecodeBytes(payload, format))
+	case "base32":
+		content = fmt.Sprintf("base64.b32decode('%s')", DecodeBytes(payload, format))
+	case "base64":
+		content = fmt.Sprintf("base64.b64decode('%s')", DecodeBytes(payload, format))
+	default:
+		content = fmt.Sprintf("'%s'", strings.Replace(DecodeBytes(payload, "ascii"), "'", "\\'", -1))
+	}
+
+	if isClient {
+		return fmt.Sprintf("p.send(%s)\n", content)
+	} else {
+		return fmt.Sprintf("p.recvuntil(%s)\n", content)
+	}
+}
diff --git a/frontend/src/backend.js b/frontend/src/backend.js
index c7abd80..1b2d8d2 100644
--- a/frontend/src/backend.js
+++ b/frontend/src/backend.js
@@ -25,6 +25,30 @@ async function json(method, url, data, json, headers) {
     }
 }
 
+async function download(url, headers) {
+
+    const options = {
+        mode: "cors",
+        cache: "no-cache",
+        credentials: "same-origin",
+        headers: headers || {},
+        redirect: "follow",
+        referrerPolicy: "no-referrer",
+    };
+    const response = await fetch(url, options);
+    const result = {
+        statusCode: response.status,
+        status: `${response.status} ${response.statusText}`,
+        blob: await response.blob()
+    };
+
+    if (response.status >= 200 && response.status < 300) {
+        return result;
+    } else {
+        return Promise.reject(result);
+    }
+}
+
 const backend = {
     get: (url = "", headers = null) =>
         json("GET", url, null, null, headers),
@@ -35,7 +59,9 @@ const backend = {
     delete: (url = "", data = null, headers = null) =>
         json("DELETE", url, null, data, headers),
     postFile: (url = "", data = null, headers = {}) =>
-        json("POST", url, data, null, headers)
+        json("POST", url, data, null, headers),
+    download: (url = "", headers = null) =>
+        download(url, headers)
 };
 
 export default backend;
diff --git a/frontend/src/components/ConnectionContent.js b/frontend/src/components/ConnectionContent.js
index 2e4f2fd..b09dcf3 100644
--- a/frontend/src/components/ConnectionContent.js
+++ b/frontend/src/components/ConnectionContent.js
@@ -7,7 +7,8 @@ import ButtonField from "./fields/ButtonField";
 import ChoiceField from "./fields/ChoiceField";
 import DOMPurify from 'dompurify';
 import ReactJson from 'react-json-view'
-import {getHeaderValue} from "../utils";
+import {downloadBlob, getHeaderValue} from "../utils";
+import log from "../log";
 
 const classNames = require('classnames');
 
@@ -92,7 +93,7 @@ class ConnectionContent extends Component {
                 if (contentType && contentType.includes("application/json")) {
                     try {
                         const json = JSON.parse(m.body);
-                        body = <ReactJson src={json} theme="grayscale" collapsed={false} displayDataTypes={false} />;
+                        body = <ReactJson src={json} theme="grayscale" collapsed={false} displayDataTypes={false}/>;
                     } catch (e) {
                         console.log(e);
                     }
@@ -126,7 +127,7 @@ class ConnectionContent extends Component {
                             messageActionDialog: <MessageAction actionName={actionName} actionValue={actionValue}
                                                                 onHide={() => this.setState({messageActionDialog: null})}/>
                         });
-                    }} />
+                    }}/>
                 );
             case "http-response":
                 const contentType = getHeaderValue(m, "Content-Type");
@@ -142,7 +143,7 @@ class ConnectionContent extends Component {
                         }
                         w.document.body.innerHTML = DOMPurify.sanitize(m.body);
                         w.focus();
-                    }} />;
+                    }}/>;
                 }
                 break;
             default:
@@ -150,6 +151,12 @@ class ConnectionContent extends Component {
         }
     };
 
+    downloadStreamRaw = (value) => {
+        backend.download(`/api/streams/${this.props.connection.id}/download?format=${this.state.format}&type=${value}`)
+            .then(res => downloadBlob(res.blob, `${this.props.connection.id}-${value}-${this.state.format}.txt`))
+            .catch(_ => log.error("Failed to download stream messages"));
+    };
+
     closeRenderWindow = () => {
         if (this.state.renderWindow) {
             this.state.renderWindow.close();
@@ -157,7 +164,8 @@ class ConnectionContent extends Component {
     };
 
     render() {
-        let content = this.state.connectionContent;
+        const conn = this.props.connection;
+        const content = this.state.connectionContent;
 
         if (content == null) {
             return <div>select a connection to view</div>;
@@ -165,7 +173,7 @@ class ConnectionContent extends Component {
 
         let payload = content.map((c, i) =>
             <div key={`content-${i}`}
-                 className={classNames("connection-message", c.from_client ? "from-client" : "from-server")}>
+                 className={classNames("connection-message", c["from_client"] ? "from-client" : "from-server")}>
                 <div className="connection-message-header container-fluid">
                     <div className="row">
                         <div className="connection-message-info col">
@@ -175,9 +183,9 @@ class ConnectionContent extends Component {
                         <div className="connection-message-actions col-auto">{this.connectionsActions(c)}</div>
                     </div>
                 </div>
-                <div className="connection-message-label">{c.from_client ? "client" : "server"}</div>
+                <div className="connection-message-label">{c["from_client"] ? "client" : "server"}</div>
                 <div
-                    className={classNames("message-content", this.state.decoded ? "message-parsed" : "message-original")}>
+                    className="message-content">
                     {this.state.tryParse && this.state.format === "default" ? this.tryParseConnectionMessage(c) : c.content}
                 </div>
             </div>
@@ -188,20 +196,20 @@ class ConnectionContent extends Component {
                 <div className="connection-content-header container-fluid">
                     <Row>
                         <div className="header-info col">
-                            <span><strong>flow</strong>: {this.props.connection.ip_src}:{this.props.connection.port_src} -> {this.props.connection.ip_dst}:{this.props.connection.port_dst}</span>
-                            <span> | <strong>timestamp</strong>: {this.props.connection.started_at}</span>
+                            <span><strong>flow</strong>: {conn["ip_src"]}:{conn["port_src"]} -> {conn["ip_dst"]}:{conn["port_dst"]}</span>
+                            <span> | <strong>timestamp</strong>: {conn["started_at"]}</span>
                         </div>
                         <div className="header-actions col-auto">
                             <ChoiceField name="format" inline small onlyName
                                          keys={["default", "hex", "hexdump", "base32", "base64", "ascii", "binary", "decimal", "octal"]}
                                          values={["plain", "hex", "hexdump", "base32", "base64", "ascii", "binary", "decimal", "octal"]}
-                                         onChange={this.setFormat} value={this.state.value} />
+                                         onChange={this.setFormat} value={this.state.value}/>
 
-                             <ChoiceField name="view_as" inline small onlyName keys={["default"]} values={["default"]} />
+                            <ChoiceField name="view_as" inline small onlyName keys={["default"]} values={["default"]}/>
 
-                            <ChoiceField name="download_as" inline small onlyName
-                                         keys={["nl_separated", "only_client", "only_server"]}
-                                         values={["nl_separated", "only_client", "only_server"]} />
+                            <ChoiceField name="download_as" inline small onlyName onChange={this.downloadStreamRaw}
+                                         keys={["nl_separated", "only_client", "only_server", "pwntools"]}
+                                         values={["nl_separated", "only_client", "only_server", "pwntools"]}/>
                         </div>
                     </Row>
                 </div>
diff --git a/frontend/src/utils.js b/frontend/src/utils.js
index fb0e5d9..aacc625 100644
--- a/frontend/src/utils.js
+++ b/frontend/src/utils.js
@@ -118,3 +118,14 @@ export function getHeaderValue(request, key) {
     }
     return undefined;
 }
+
+export function downloadBlob(blob, fileName) {
+    const url = window.URL.createObjectURL(blob);
+    const a = document.createElement("a");
+    a.style.display = "none";
+    a.href = url;
+    a.download = fileName;
+    document.body.appendChild(a);
+    a.click();
+    window.URL.revokeObjectURL(url);
+}