//.............................................................................
//
// Sniffer session
//

import "io_base.jncx"
import "io_pcap.jncx"
import "doc_Plugin.jnc"
import "gui_Dialog.jnc"

//.............................................................................

class SnifferSession: doc.Session
{
protected:
	enum State
	{
		Closed,
		Capturing,
	}

	enum ActionId
	{
		Capture,
		_Count,
	}

	enum IconId
	{
		Capture,
		StopCapture,
		_Count,
	}

	enum StatusPaneId
	{
		State,
		_Count,
	}

protected:
	// pcap

	State bindable m_state;
	io.PCap m_pcap;	
	io.IpDefragmenter m_ipDefragmenter;
	
	// toolbar, status bar & menu
	
	gui.Action* m_actionTable [ActionId._Count];
	gui.Icon* m_iconTable [IconId._Count];
	gui.StatusPane* m_statusPaneTable [StatusPaneId._Count];
	gui.ComboBox* m_deviceCombo;
	gui.ComboBox* m_filterCombo;
	gui.FileDlg* m_fileDlg;
	
	// property grid
	
	gui.EnumProperty* m_deviceProp;
	gui.BoolProperty* m_promisciousProp;
	gui.StringProperty* m_filterProp;
	gui.IntProperty* m_ipFragmentLimitProp;
	gui.IntProperty* m_ipFragmentTimeoutProp;

public:
	construct (doc.PluginHost* pluginHost);

	bool capture () throws;
	
	bool capture (
		io.PCapDeviceDesc const* device,
		char const* filterString
		) throws;		

	bool openCapFile (
		char const* fileName,
		char const* filterString
		) throws;

	stopCapture ();

	override void load (doc.Storage* storage);
	override void save (doc.Storage* storage);
	override void updateProperties ();
	override bool applyProperties () throws;
	override void restoreDefaultProperties ();

	override size_t transmit (
		void const* p,
		size_t size
		) throws;

protected:
	onCapture ();
	onFileDlgCompleted (gui.StdDlgButton button);
	onSetFilter ();
	onPCapEvent (io.PCapEventParams const* params);

	bool applyFilter () throws;

	processPacket (
		void const* p,
		size_t size
		);

	createUi ();
	reactor m_uiReactor ();
}

//. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

SnifferSession.construct (doc.PluginHost* pluginHost)
{
	basetype.construct (pluginHost);

	pluginHost.m_log.m_representerFunc = representSnifferLog;

	m_state = State.Closed;
	m_pcap.m_onPCapEvent += onPCapEvent @ pluginHost.m_mainThreadScheduler;

	createUi ();
}

SnifferSession.createUi ()
{
	// toolbar

	m_pluginHost.m_toolBar.addLabel ("Device:");
	m_deviceCombo = io.createPCapDeviceComboBox (m_pluginHost.m_toolBar);
	m_deviceCombo.m_onEnter += onCapture;

	m_pluginHost.m_toolBar.addLabel ("Filter:");
	m_filterCombo = m_pluginHost.m_toolBar.addComboBox (-1);
	m_filterCombo.m_isEditable = true;
	m_filterCombo.m_onEnter += onSetFilter;

	m_iconTable [IconId.Capture] = new gui.Icon ("images/capture.png");
	m_iconTable [IconId.StopCapture] = new gui.Icon ("images/pause.png");

	m_actionTable [ActionId.Capture] = createAction ("Capture", m_iconTable [IconId.Capture]);
	m_actionTable [ActionId.Capture].m_onTriggered += onCapture;

	// status bar
	
	m_statusPaneTable [StatusPaneId.State] = m_pluginHost.m_statusBar.addPane (, 80);

	// property grid

	m_deviceProp = io.createPCapDeviceProperty (
		m_pluginHost.m_propertyGrid,,
		"Device",
		"Specify the device to capture on"
		);

	m_promisciousProp = m_pluginHost.m_propertyGrid.createBoolProperty (
		"Promiscious Mode",
		"Put the capturing interface into promiscious mode"
		);

	m_filterProp = m_pluginHost.m_propertyGrid.createStringProperty (
		"Capture Filter",
		"Capture only the packets which have passed this primary filter"
		);

	m_ipFragmentLimitProp  = m_pluginHost.m_propertyGrid.createIntProperty (
		"IP Fragment Limit",
		"Specify the maximum number of IP fragments"
		);

	m_ipFragmentTimeoutProp  = m_pluginHost.m_propertyGrid.createIntProperty (
		"IP Fragment Timeout (ms)",
		"Specify the maximum delay between IP fragments"
		);

	m_uiReactor.start ();
}

SnifferSession.onCapture ()
{
	if (m_state)
		stopCapture ();
	else
		try capture ();
}

SnifferSession.onSetFilter ()
{
	if (!m_state)
		try capture ();
	else
		try applyFilter ();	
}

SnifferSession.onFileDlgCompleted (gui.StdDlgButton button)
{
	if (button != gui.StdDlgButton.Ok)
		return;
	
	char const* fileName = m_fileDlg.m_filePath;
	char const* filterString = m_filterCombo.m_editText;
	m_filterCombo.addEditTextToHistory ();

	try openCapFile (fileName, filterString);
}

bool SnifferSession.capture () throws
{
	stopCapture ();
	
	if (m_deviceCombo.m_currentIndex == -1)
		return false;

	io.PCapDeviceDesc const* device = m_deviceCombo.m_currentData;
	if (!device)
	{
		if (!m_fileDlg)
		{
			m_fileDlg = new gui.FileDlg;
			m_fileDlg.m_onCompleted += onFileDlgCompleted;			
		}
		
		m_fileDlg.m_kind = gui.FileDlgKind.Open;
		m_fileDlg.m_title = "Open Wireshark File (*.cap; *.pcap)...";
		m_fileDlg.m_filter = "Wireshark .cap files (*.cap; *.pcap);; All files (*.*)";
		m_fileDlg.show ();
		return true;
	}

	char const* filterString = m_filterCombo.m_editText;
	m_filterCombo.addEditTextToHistory ();

	return try capture (device, filterString);
}

bool SnifferSession.capture (
	io.PCapDeviceDesc const* device,
	char const* filterString
	) throws
{
	stopCapture ();
		
	writeLogString (SnifferLogRecordCode.StartingCapture, filterString);
	
	bool result = try m_pcap.openDevice (
		device.m_name, 
		filterString,
		m_promisciousProp.m_value
		);
		
	if (!result)
	{
		writeLogLastError (SnifferLogRecordCode.CaptureError);
		return false;
	}

	m_state = State.Capturing;
	writeLog (SnifferLogRecordCode.CaptureStarted);
	return true;
}

bool SnifferSession.openCapFile (
	char const* fileName,
	char const* filterString
	) throws
{
	writeLogString (SnifferLogRecordCode.StartingCapture, filterString);
	
	bool result = try m_pcap.openFile (
		fileName, 
		filterString
		);
		
	if (!result)
	{
		writeLogLastError (SnifferLogRecordCode.CaptureError);
		return false;
	}

	m_state = State.Capturing;
	writeLog (SnifferLogRecordCode.CaptureStarted);
	return true;
}

SnifferSession.stopCapture ()
{
	if (!m_state)
		return;

	writeLog (SnifferLogRecordCode.CaptureStopped);
	m_pcap.close ();
	m_state = State.Closed;
}

bool SnifferSession.applyFilter () throws
{
	char const* filter = m_filterCombo.m_editText;
	m_filterCombo.addEditTextToHistory ();

	if (!m_state.Capturing)
		return true;

	writeLogString (SnifferLogRecordCode.ChangingFilter, filter);
	bool result = try m_pcap.setFilter (filter);
	if (!result)
	{
		writeLogLastError (SnifferLogRecordCode.FilterError);
		return false;
	}
	
	writeLog (SnifferLogRecordCode.FilterChanged);
	return true;
}

void SnifferSession.load (doc.Storage* storage)
{
	m_deviceCombo.m_currentIndex = storage.readInt ("deviceIdx");
	m_filterCombo.m_editText = storage.readString ("captureFilter");
	m_promisciousProp.m_value = storage.readBool ("promiscious");
	m_ipFragmentLimitProp.m_value = storage.readInt ("ipFragmentLimit");
	m_ipFragmentTimeoutProp.m_value = storage.readInt ("ipFragmentTimeout");
}

void SnifferSession.save (doc.Storage* storage)
{
	storage.writeInt ("deviceIdx", (int) m_deviceCombo.m_currentIndex);
	storage.writeString ("captureFilter", m_filterCombo.m_editText);
	storage.writeBool ("promiscious", m_promisciousProp.m_value);
	storage.writeInt ("ipFragmentLimit", m_ipFragmentLimitProp.m_value);
	storage.writeInt ("ipFragmentTimeout", m_ipFragmentTimeoutProp.m_value);
}

void SnifferSession.updateProperties ()
{
	m_deviceProp.m_currentIndex = m_deviceCombo.m_currentIndex;
	m_filterProp.m_value = m_filterCombo.m_editText;
}

bool SnifferSession.applyProperties () throws
{
	m_deviceCombo.m_currentIndex = m_deviceProp.m_currentIndex;
	m_filterCombo.m_editText = m_filterProp.m_value;
	m_ipDefragmenter.m_fragmentLimit = m_ipFragmentLimitProp.m_value;
	m_ipDefragmenter.m_timeout = m_ipFragmentTimeoutProp.m_value;
	
	if (m_state)
		try applyFilter ();
	
	return true;
}

void SnifferSession.restoreDefaultProperties ()
{
	m_deviceProp.m_currentIndex = 0;
	m_promisciousProp.m_value = false;
	m_filterProp.m_value = "tcp and port 80";
	m_ipFragmentLimitProp.m_value = io.IpDefragmenter.DefIpFragmentLimit;
	m_ipFragmentTimeoutProp.m_value = io.IpDefragmenter.DefIpFragmentTimeout;
}

size_t SnifferSession.transmit (
	void const* p,
	size_t size
	) throws
{
	if (!m_pcap)
		return -1;

	size_t result = try m_pcap.write (p, size);
	if (result != -1)
		writeLog (SnifferLogRecordCode.InjectedPacket, p, result);

	return result;
}

SnifferSession.onPCapEvent (io.PCapEventParams const* params)
{
	if (params.m_syncId != m_pcap.m_syncId)
		return;

	switch (params.m_eventKind)
	{
	case io.PCapEventKind.ReadyRead:
		static char buffer [1600]; // static is OK -- we read in main thread
		size_t result = try m_pcap.read (buffer, sizeof (buffer));
		if (result != -1)
			processPacket (buffer, result);
		
		break;
	
	case io.PCapEventKind.Eof:
		stopCapture ();
		break;
	}
}

SnifferSession.processPacket (
	void const* p,
	size_t size
	)
{
	writeLog (SnifferLogRecordCode.Packet, p, size);

	// defragment IPv4
	
	if (size < sizeof (io.EthernetHdr) + sizeof (io.IpHdr))
		return;
	
	io.EthernetHdr const* ethernetHdr = (io.EthernetHdr const*) p;
	if (ethernetHdr.m_type != io.EthernetType.Ip)
		return;

	p += sizeof (io.EthernetHdr);
	size -= sizeof (io.EthernetHdr);
	
	io.IpHdr const* ipHdr = (io.IpHdr const*) p;
	if ((ipHdr.m_flags & io.IpFlags.Mf) || ipHdr.m_fragmentOffset)
	{
		ipHdr = m_ipDefragmenter.defragment (ipHdr, size);
		if (ipHdr)
		{
			size = dynamic sizeof (ipHdr);
			writeLog (SnifferLogRecordCode.DefragmentedIp, ipHdr, size);
		}
	}
}

reactor SnifferSession.m_uiReactor ()
{
	m_title = $"PCap $(m_deviceCombo.m_currentText)";
	m_isTransmitEnabled = m_state == State.Capturing;
	m_actionTable [ActionId.Capture].m_text = m_state ? "Stop Capture" : "Capture";
	m_actionTable [ActionId.Capture].m_icon = m_iconTable [m_state ? IconId.StopCapture : IconId.Capture];
	m_statusPaneTable [StatusPaneId.State].m_text = m_state ? "Capturing" : "Closed";
}

//.............................................................................
