################################################################################
## WebSockets Protocol Analyzer Spicy Script
## Jennifer Gates
## August 2017
##
## This module parses HTTP GET request messages looking for the GET string and 
## the  “Sec-WebSocket-Version” header. If both found, follow-on packets are 
## considered WebSockets protocol packets and parsed accordingly.  HTTP responses
## starting with “HTTP/1.1 101 Switching Proto” and the “Sec-WebSocket-Accept” 
## header. If both found, follow-on packets are considered WebSockets protocol 
## packets and parsed accordingly.
################################################################################
module WS_HANDSHAKE;

import Spicy;
import "HILTI-C" void Hilti::terminate();

const BeforeColon = /[^:]+/;
const Colon = /: /;
const DataValue = /[^\x0d\x0a]+/;
const LineEnd = /\x0d\x0a/;
const DataEnd = b"\x53\x50\x49\x43\x59\x6a\x6c\x67";

global gtotaldata: bytes =b"";


# function to add a distinguishing byte at the end of the ws packets 
# during parsing to allow list to properly terminate
bytes terminate(b: bytes) {
	return b + DataEnd;
}

# primary unit to store client originator Websocket handshake traffic  
# uses sink to parse follow on websocket protocol data packets
export type WS_Handshake = unit {

    get    : /^(GET|get|Get) /;
	dvalue : DataValue;
	headers : list<Header>;
	end_of_hdrs : /\x0d\x0a\x0d\x0a/;

	ws_data : bytes &eod &convert=terminate($$) &transient  -> self.hs_sub;

	on %init {
		self.hs_sub.connect(new HS_Sub(self));
	}

	var hs_sub: sink;
	var messages: list<WS_Message>;
	var totaldata: bytes;

	# stops parsing if HTTP GET request is not initiating Websocket handshake
	# joins all headers into a string and searches for Sec-WebSocket-Version header
	on headers {
		local joinedheaders: bytes;
		joinedheaders = b".".join(self.headers);
		joinedheaders = joinedheaders.lower();
		if (|joinedheaders.match(/.*sec-websocket-version.*/)| == 0){
			Hilti::terminate();
		}
	}

	# combines data fields from all websockset packets into one string
    on %done { 
		self.totaldata = gtotaldata;
    	gtotaldata = b" ";
		self.hs_sub.close();
	}
};

# sub unit to put originator websocket protocol packets into a list
export type HS_Sub = unit(handshake: WS_Handshake) {
	ws_msgs : list<WS_Message> {
		handshake.messages = self.ws_msgs;
	}	
		: DataEnd;
};

# primary unit to store server/responder Websocket handshake traffic 
# uses sink to parse follow on websocket protocol data packets 
export type WS_Handshake_Success = unit {

    success     : /^(HTTP|http)\/1.1 101 Switching Proto/;
    dvalue      : DataValue;
    svrheaders     : list<Header>;
    end_of_hdrs : /\x0d\x0a\x0d\x0a/;
    wss_data : bytes &eod &convert=terminate($$) &transient -> self.hss_sub;

	on %init {
		self.hss_sub.connect(new HSS_Sub(self));
    }

    var hss_sub: sink;
	var smessages: list<WS_Message>;
	var totaldata: bytes;

	# stops parsing if HTTP reply is not responding to Websocket handshake
	# joins all headers into a string and searches for Sec-WebSocket-Accept header
	on svrheaders {
		local sjoinedheaders: bytes;
		sjoinedheaders = b".".join(self.svrheaders);
		sjoinedheaders = sjoinedheaders.lower();
		if (|sjoinedheaders.match(/.*sec-websocket-accept.*/)| == 0){
			Hilti::terminate();
		}
	}

	# combines data fields from all websockset packets into one string
    on %done {
		self.totaldata = gtotaldata;
		gtotaldata = b" ";
		self.hss_sub.close();
    }
};

# sub unit to put responder websocket protocol packets into a list
export type HSS_Sub = unit (shandshake: WS_Handshake_Success) {
        wss_msgs : list<WS_Message>{
		shandshake.smessages = self.wss_msgs;
	}
       : DataEnd;
};

# unit to parse http headers in handshake packets
export type Header = unit {
		    : LineEnd;
	name	: BeforeColon;
		    : Colon;
	value	: DataValue;
};

# sub unit to parse WS protocol fields from websockets packet list
export type WS_Message = unit {
	# start with first 2 bytes as a bitfield to get access to bit level flags
    first2B: bitfield(16) {
        fin: 0;
        rsv1: 1;
        rsv2: 2;
        rsv3: 3;
        op: 4..7;
        mask: 8;
        pay1: 9..15;
    } &bitorder = Spicy::BitOrder::MSB0;
	
	# determine payload length and start point for next field
	switch ( self.first2B.pay1 ) {
		126 -> extpay: ExtPay;
		127 -> extpayc: ExtPayC;
		* -> : void;
	};
	
	# parse out masking key if payload is masked (from client)
	switch ( self.first2B.mask ) {
		1 -> maskkey: bytes &length=4;
		0 -> : void;
	};
	
	# grab the payload data. Leave unmasking to bro script
    switch ( self.first2B.pay1 ) {
        126 -> data: bytes &length=self.extpay.pay2 ;
        127 -> data: bytes &length=self.extpayc.pay3 ;
        * -> data: bytes &length=self.first2B.pay1;
    };

    # append data from this packet to connection's data string
	on data {
		gtotaldata += self.data;
	}

};


# define this new type which parses out the extended payload length field 
export type ExtPay = unit {
	pay2: uint<16>;
};

# define this new type which parses out the extended continued payload length field
export type ExtPayC = unit {
	pay3: uint<64>;
};