Add functions to write TLS key-logs to PcapNG files

pcapgo/ngwrite.go

@@ -391,6 +391,94 @@ func (w *NgWriter) WritePacket(ci gopacket.CaptureInfo, data []byte) error {
 	return err
 }
 
+// WriteCustomBlock writes an arbitrary block to a PcapNG file. Standard
+// block types are specified in section 10.1 of the PcapNG specification.
+// For example, an interface statistics block has a block type of 0x00000005.
+// According to the spec, all blocks must be padded to a 32-bit boundary.
+// If len(body) is not a multiple of 4, then zero-padding will be added to
+// the end of the body to make it a multiple of 4.
+func (w *NgWriter) WriteCustomBlock(blockType uint32, body []byte) error {
+	bodyLength := len(body)
+	padding := (4 - bodyLength&3) & 3
+	totalLength := bodyLength + padding + 12
+
+	err := binary.Write(w.w, binary.LittleEndian, blockType)
+	if err != nil {
+		return err
+	}
+
+	err = binary.Write(w.w, binary.LittleEndian, uint32(totalLength))
+	if err != nil {
+		return err
+	}
+
+	_, err = w.w.Write(body)
+	if err != nil {
+		return err
+	}
+
+	_, err = w.w.Write(make([]byte, padding))
+	if err != nil {
+		return err
+	}
+
+	err = binary.Write(w.w, binary.LittleEndian, uint32(totalLength))
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// WriteTLSKeyLog writes a TLS key log block to a PcapNG file. More
+// technically, it writes a decryption secrets block, with a secrets type of
+// "TLS Key Log" (0x544c534b). This expects a key log in the format
+// generated by crypto/tls's tls.Config.KeyLogWriter (many other tools also
+// generate this format). According to the spec, secrets blocks should
+// appear before any packets that they decrypt.
+func (w *NgWriter) WriteTLSKeyLog(kl []byte) error {
+	klLength := len(kl)
+	padding := (4 - klLength&3) & 3
+	totalLength := klLength + padding + 20
+
+	err := binary.Write(w.w, binary.LittleEndian, uint32(ngBlockTypeDecryptionSecrets))
+	if err != nil {
+		return err
+	}
+
+	err = binary.Write(w.w, binary.LittleEndian, uint32(totalLength))
+	if err != nil {
+		return err
+	}
+
+	err = binary.Write(w.w, binary.LittleEndian, uint32(0x544c534b)) // NSS key log format
+	if err != nil {
+		return err
+	}
+
+	err = binary.Write(w.w, binary.LittleEndian, uint32(klLength))
+	if err != nil {
+		return err
+	}
+
+	_, err = w.w.Write(kl)
+	if err != nil {
+		return err
+	}
+
+	_, err = w.w.Write(make([]byte, padding))
+	if err != nil {
+		return err
+	}
+
+	err = binary.Write(w.w, binary.LittleEndian, uint32(totalLength))
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // Flush writes out buffered data to the storage media. Must be called before closing the underlying file.
 func (w *NgWriter) Flush() error {
 	return w.w.Flush()

pcapgo/ngwrite_test.go

@@ -7,6 +7,7 @@
 package pcapgo
 
 import (
+	"bufio"
 	"bytes"
 	"testing"
 	"time"
@@ -214,6 +215,118 @@ func TestNgWriteComplex(t *testing.T) {
 	ngRunFileReadTest(test, "", false, t)
 }
 
+func TestWriteTLSKeyLog(t *testing.T) {
+	// Create a buffer to capture the output
+	buffer := new(bytes.Buffer)
+
+	// Create a new NgWriter with the buffer
+	w := &NgWriter{
+		w: bufio.NewWriter(buffer),
+	}
+
+	// Call the WriteTLSKeyLog function with some test data
+	kl := []byte("test key log!")
+	w.WriteTLSKeyLog(kl)
+
+	// Flush the buffer and capture the result
+	w.Flush()
+	result := buffer.Bytes()
+
+	if len(result) != 36 {
+		t.Fatalf("Expected 36 bytes, got %d", len(result))
+	}
+	if !bytes.Equal(result[0:4], []byte{0x0A, 0x00, 0x00, 0x00}) {
+		t.Fatalf("Unexpected block type %x", result[0:4])
+	}
+	if !bytes.Equal(result[4:8], []byte{36, 0, 0, 0}) {
+		t.Fatalf("Unexpected value in 1st length field %v", result[4:8])
+	}
+	if !bytes.Equal(result[8:12], []byte{0x4b, 0x53, 0x4c, 0x54}) {
+		t.Fatalf("Unexpected key log format %x", result[8:12])
+	}
+	if !bytes.Equal(result[12:16], []byte{13, 0, 0, 0}) {
+		t.Fatalf("Unexpected value in key log length field %x", result[12:16])
+	}
+	if !bytes.Equal(result[16:29], []byte("test key log!")) {
+		t.Fatalf(`Unexpected key log data "%s"`, result[16:29])
+	}
+	if !bytes.Equal(result[29:32], []byte{0, 0, 0}) {
+		t.Fatalf("Expected zero-padding, got %x", result[29:32])
+	}
+	if !bytes.Equal(result[32:36], []byte{36, 0, 0, 0}) {
+		t.Fatalf("Unexpected value in 2nd length field %v", result[32:36])
+	}
+}
+
+func TestWriteCustomBlock(t *testing.T) {
+	// Create a buffer to capture the output
+	buffer := new(bytes.Buffer)
+
+	// Create a new NgWriter with the buffer
+	w := &NgWriter{
+		w: bufio.NewWriter(buffer),
+	}
+
+	// Write a custom block
+	blockType := uint32(0xDEADBEEF)
+	body := []byte("test body")
+	w.WriteCustomBlock(blockType, body)
+
+	// Flush the buffer and capture the result
+	w.Flush()
+	result := buffer.Bytes()
+
+	if len(result) != 24 {
+		t.Fatalf("Expected 24 bytes, got %d", len(result))
+	}
+	if !bytes.Equal(result[0:4], []byte{0xEF, 0xBE, 0xAD, 0xDE}) {
+		t.Fatalf("Unexpected block type %x", result[0:4])
+	}
+	if !bytes.Equal(result[4:8], []byte{24, 0, 0, 0}) {
+		t.Fatalf("Unexpected value in 1st length field %v", result[4:8])
+	}
+	if !bytes.Equal(result[8:17], []byte("test body")) {
+		t.Fatalf(`Unexpected body "%s"`, result[8:17])
+	}
+	if !bytes.Equal(result[17:20], []byte{0, 0, 0}) {
+		t.Fatalf("Expected zero-padding, got %x", result[17:20])
+	}
+	if !bytes.Equal(result[20:24], []byte{24, 0, 0, 0}) {
+		t.Fatalf("Unexpected value in 2nd length field %v", result[20:24])
+	}
+}
+
+func TestWriteCustomBlock_EmptyBody(t *testing.T) {
+	// Create a buffer to capture the output
+	buffer := new(bytes.Buffer)
+
+	// Create a new NgWriter with the buffer
+	w := &NgWriter{
+		w: bufio.NewWriter(buffer),
+	}
+
+	// Write a custom block with an empty body
+	blockType := uint32(0xABADFACE)
+	w.WriteCustomBlock(blockType, nil)
+
+	// Flush the buffer and capture the result
+	w.Flush()
+	result := buffer.Bytes()
+
+	if len(result) != 12 {
+		t.Fatalf("Expected 16 bytes, got %d", len(result))
+	}
+	if !bytes.Equal(result[0:4], []byte{0xCE, 0xFA, 0xAD, 0xAB}) {
+		t.Fatalf("Unexpected block type %x", result[0:4])
+	}
+	if !bytes.Equal(result[4:8], []byte{12, 0, 0, 0}) {
+		t.Fatalf("Unexpected value in 1st length field %v", result[4:8])
+	}
+	if !bytes.Equal(result[8:12], []byte{12, 0, 0, 0}) {
+		t.Fatalf("Unexpected value in 2nd length field %v", result[12:16])
+	}
+}
+
 type ngDevNull struct{}
 
 func (w *ngDevNull) Write(p []byte) (n int, err error) {

pcapgo/pcapng.go

@@ -37,6 +37,7 @@ const (
 	ngBlockTypeSimplePacket        ngBlockType = 3          // Simple packet block
 	ngBlockTypeInterfaceStatistics ngBlockType = 5          // Interface statistics block
 	ngBlockTypeEnhancedPacket      ngBlockType = 6          // Enhanced packet block
+	ngBlockTypeDecryptionSecrets   ngBlockType = 10         // Decryption secrets block
 	ngBlockTypeSectionHeader       ngBlockType = 0x0A0D0D0A // Section header block (same in both endians)
 )