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

import "doc_Session.jnc"
import "io_base.jncx"
import "io_pcap.jncx"
import "io_IpDefragmenter.jnc"
import "ui_Dialog.jnc"
import "ui_History.jnc"
import "ui_PcapUi.jnc"
import "ui_BufferPropertySet.jnc"
import "ui_StdSessionInfoSet.jnc"
import "log_ThroughputCalc.jnc"
import "log_ChecksumCalc.jnc"
import "log_exportToPcap.jnc"
import "std_Buffer.jnc"

using extension ui.ComboBoxHistory;

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

class PcapSnifferSession: doc.Session {
protected:
	enum Defaults {
		Promiscious      = false,
		SnapshotSize     = 64 * 1024,
		KernelBufferSise = 0,
		BufferSize       = 1 * 1024 * 1024,
	}

	enum State {
		Closed,
		Capturing,
	}

	enum ActionId {
		ApplyFilter,
		Capture,
		OpenPcap,
		ExportPcap,
		_Count,
	}

	enum IconId {
		Capture,
		StopCapture,
		_Count,
	}

	enum StatusPaneId {
		State,
		_Count,
	}

protected:
	// pcap

	State bindable m_state;
	io.Pcap m_pcap;
	char const* m_captureFilterString;
	io.IpDefragmenter m_ipDefragmenter;
	std.Buffer m_readBuffer;
	NetSnifferLogFilter m_filter;
	uint_t m_syncId;

	// toolbar, status bar & menu

	ui.Action* m_actionTable[ActionId._Count];
	ui.Icon* m_iconTable[IconId._Count];
	ui.StatusPane* m_statusPaneTable[StatusPaneId._Count];
	ui.ComboBox* m_viewFilterCombo;
	ui.ComboBox* m_deviceCombo;
	ui.FileDlg* m_fileDlg;

	// property grid

	ui.EnumProperty* m_deviceProp;
	ui.BoolProperty* m_promisciousProp;
	ui.StringProperty* m_captureFilterProp;
	ui.StringProperty* m_viewFilterProp;
	ui.IntProperty* m_ipFragmentLimitProp;
	ui.IntProperty* m_ipFragmentTimeoutProp;
	ui.IntProperty* m_snapshotSizeProp;
	ui.IntProperty* m_kernelBufferSizeProp;
	ui.BufferPropertySet m_bufferPropertySet;

	// information grid

	ui.StdSessionInfoSet m_stdSessionInfoSet;
	ui.ChecksumInfoSet m_checksumInfoSet;

	// ui reactor

	reactor m_uiReactor;

public:
	construct(doc.PluginHost* pluginHost);

	bool errorcode capture();

	bool errorcode capture(
		io.PcapDeviceDesc const* device,
		char const* filterString
	);

	bool errorcode openPcapFile(
		char const* fileName,
		char const* filterString
	);

	void stopCapture();

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

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

protected:
	void onCapture() {
		try(m_state ? stopCapture() : capture());
	}

	void onOpenPcap();
	void onExportPcap();

	void onSetViewFilter() {
		try applyViewFilter(true);
	}

	void onDeviceComboEnter() {
		try capture();
	}

	void onFileDlgCompleted(ui.StdDlgButton button);

	void onPcapEvent(
		uint_t syncId,
		io.PcapEvents triggeredEvents
	);

	bool errorcode applyCaptureFilter();
	bool errorcode applyViewFilter(bool rebuildLog = false);

	void processPacket(
		uint64_t timestamp,
		void const* p,
		size_t size
	);

	void ensureFileDlg();
	void waitPcapEvents();
	void createUi();
}

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

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

	createUi();

	pluginHost.m_log.addSelectionProcessor(log.ChecksumCalcRangeProcessor.factory~(m_checksumInfoSet));
	pluginHost.m_log.addFilter(m_filter);
}

void PcapSnifferSession.createUi() {
	// toolbar

	m_pluginHost.m_toolBar.addLabel("Filter:");
	m_viewFilterCombo = m_pluginHost.m_toolBar.addComboBox(-1);
	m_viewFilterCombo.m_isEditable = true;
	m_viewFilterCombo.m_onEnter += onSetViewFilter;
	m_viewFilterCombo.m_placeholderText = "Enter a pcap-filter...";
	m_viewFilterCombo.m_toolTipText =
		"View only packets which pass this filter\n"
		"For a complete syntax reference, google 'pcap-filter'";

	m_actionTable[ActionId.ApplyFilter] = createAction("Apply filter", "images/apply.png");
	m_actionTable[ActionId.ApplyFilter].m_onTriggered += onSetViewFilter;

	m_pluginHost.m_toolBar.addLabel("Device:");
	m_deviceCombo = ui.createPcapDeviceComboBox(m_pluginHost.m_toolBar);
	m_deviceCombo.m_onEnter += onDeviceComboEnter;

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

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

	m_pluginHost.m_menu.addSeparator();
	m_pluginHost.m_toolBar.addSeparator();

	m_actionTable[ActionId.OpenPcap] = createAction("Open Pcap file (*.pcap)", "images/open-pcap.png");
	m_actionTable[ActionId.OpenPcap].m_onTriggered += onOpenPcap;
	m_actionTable[ActionId.ExportPcap] = createAction("Export Pcap file (*.pcap)", "images/save-pcap.png");
	m_actionTable[ActionId.ExportPcap].m_onTriggered += onExportPcap;

	// status bar

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

	// property grid

	ui.GroupProperty* group = m_pluginHost.m_propertyGrid.createGroupProperty(
		"Connection",
		"Pcap-specific connection settings"
	);

	m_deviceProp = ui.createPcapDeviceProperty(
		m_pluginHost.m_propertyGrid,
		group,,
		"Device",
		"Specify the device to capture on"
	);

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

	m_captureFilterProp = m_pluginHost.m_propertyGrid.createStringProperty(
		group,,
		"Capture filter",
		"Enter a pcap-filter...",
		"Capture only packets which pass this filter\n"
		"For a complete syntax reference, google 'pcap-filter'"
	);

	m_viewFilterProp = m_pluginHost.m_propertyGrid.createStringProperty(
		group,,
		"View filter",
		"Enter a pcap-filter...",
		"View only packets which pass this filter\n"
		"For a complete syntax reference, google 'pcap-filter'"
	);

	group = m_pluginHost.m_propertyGrid.createGroupProperty(
		"IP defragmenter",
		"IP defragmentation settings"
	);

	m_ipFragmentLimitProp  = m_pluginHost.m_propertyGrid.createIntProperty(
		group,,
		"IP fragment limit",
		"Specify the maximum number of IP fragments"
	);

	m_ipFragmentTimeoutProp  = m_pluginHost.m_propertyGrid.createIntProperty(
		group,,
		"IP fragment timeout (ms)",
		"Specify the maximum delay between IP fragments"
	);

	group = m_pluginHost.m_propertyGrid.createGroupProperty(
		"Buffering",
		"Advanced buffering options affecting throughtput and compatibility"
	);

	m_snapshotSizeProp = m_pluginHost.m_propertyGrid.createIntProperty(
		group,,
		"Snapshot size",
		"Specify the Pcap snapshot size"
	);

	m_snapshotSizeProp.m_spinBoxStep = 1024;

	m_kernelBufferSizeProp = m_pluginHost.m_propertyGrid.createIntProperty(
		group,,
		"Kernel buffer size",
		"Specify the Pcap kernel buffer size (0 = default)"
	);

	m_kernelBufferSizeProp.m_spinBoxStep = 1024;

	m_bufferPropertySet.createProperties(
		m_pluginHost.m_propertyGrid,
		group,,
		ui.BufferPropertySet.PropertyId.ReadBufferSize |
		ui.BufferPropertySet.PropertyId.NotifyReadBufferFull
	);

	m_bufferPropertySet.m_defaultReadBufferSize = Defaults.BufferSize;

	// information grid

	ui.InformationGroup* infoGroup = m_pluginHost.m_informationGrid.createGroup("Pcap sniffer");

	m_stdSessionInfoSet.create(
		m_pluginHost.m_txRxStats,
		m_pluginHost.m_informationGrid,
		infoGroup,,
		ui.StdSessionInfoSet.InfoId.SessionTime |
		ui.StdSessionInfoSet.InfoId.TotalSize |
		ui.StdSessionInfoSet.InfoId.TotalThroughput
	);

	m_checksumInfoSet.createInformationValuesGrouped(m_pluginHost.m_informationGrid);

	m_pluginHost.m_propertyGrid.resizeNameColumnToContents();
	m_pluginHost.m_informationGrid.resizeNameColumnToContents();

	m_uiReactor.start();
}

void PcapSnifferSession.onOpenPcap() {
	ensureFileDlg();

	m_fileDlg.m_kind = ui.FileDlgKind.Open;
	m_fileDlg.m_title = "Open Pcap File...";
	m_fileDlg.m_filter = "Pcap files (*.pcap; *.cap);; All files (*.*)";
	m_fileDlg.show();
}

void PcapSnifferSession.onExportPcap() {
	ensureFileDlg();

	m_fileDlg.m_kind = ui.FileDlgKind.Save;
	m_fileDlg.m_title = "Export Pcap File...";
	m_fileDlg.m_filter = "Pcap files (*.pcap; *.cap);; All files (*.*)";
	m_fileDlg.show();
}

size_t exportLogRecordToPcap(
	log.Record const* record,
	io.MappedFile* pcapFile,
	uint64_t pcapOffset
) {
	void const* p;
	size_t size;

	switch (record.m_code) {
	case NetSnifferLogRecordCode.EthernetPacket:
	case NetSnifferLogRecordCode.InjectedEthernetPacket:
		p = record + 1;
		size = record.m_dataSize;
		break;

	case NetSnifferLogRecordCode.Packet:
	case NetSnifferLogRecordCode.InjectedPacket:
		p = record + 1;
		p += sizeof(io.PcapLinkType);
		size = record.m_dataSize - sizeof(io.PcapLinkType);
		break;

	default:
		return 0;
	}

	return log.writePcapPacket(pcapFile, pcapOffset, record.m_timestamp, p, size);
}

void PcapSnifferSession.onFileDlgCompleted(ui.StdDlgButton button) {
	if (button != ui.StdDlgButton.Ok)
		return;

	switch (m_fileDlg.m_kind) {
	case ui.FileDlgKind.Open:
		try openPcapFile(m_fileDlg.m_filePath, m_captureFilterProp.m_value);
		break;

	case ui.FileDlgKind.Save:
		bool result = try log.exportToPcap(
			m_pluginHost.m_log,
			m_fileDlg.m_filePath,
			io.PcapLinkType.Ethernet,
			m_snapshotSizeProp.m_value,
			exportLogRecordToPcap
		);

		if (!result)
			m_logWriter.writeLastError(log.StdRecordCode.Error);

		break;
	}
}

void PcapSnifferSession.ensureFileDlg() {
	if (!m_fileDlg) {
		m_fileDlg = new ui.FileDlg;
		m_fileDlg.m_onCompleted += onFileDlgCompleted;
	}
}

bool errorcode PcapSnifferSession.capture() {
	stopCapture();

	io.PcapDeviceDesc const* device = m_deviceCombo.m_currentData;
	return device ? try capture(device, m_captureFilterProp.m_value) : false;
}

bool errorcode PcapSnifferSession.capture(
	io.PcapDeviceDesc const* device,
	char const* filterString
) {
	stopCapture();
	startSession();

	m_logWriter.write(PcapLogRecordCode.StartingCapture, filterString);

	m_pcap.openDevice(device.m_name);
	m_pcap.m_isPromiscious = m_promisciousProp.m_value;
	m_pcap.m_snapshotSize = m_snapshotSizeProp.m_value;
	m_pcap.m_kernelBufferSize = m_kernelBufferSizeProp.m_value;
	m_pcap.m_readBufferSize = m_bufferPropertySet.m_readBufferSizeProp.m_value;
	m_pcap.activate(filterString);

	m_captureFilterString = filterString;
	m_state = State.Capturing;
	m_logWriter.write(PcapLogRecordCode.CaptureStarted);
	waitPcapEvents();
	m_stdSessionInfoSet.startSession(m_pluginHost.m_mainThreadScheduler);
	applyViewFilter();
	return true;

catch:
	m_logWriter.writeLastError(PcapLogRecordCode.CaptureError);
	m_pcap.close();
	return false;
}

bool errorcode PcapSnifferSession.openPcapFile(
	char const* fileName,
	char const* filterString
) {
	m_logWriter.write(PcapLogRecordCode.StartingCapture, filterString);

	m_pcap.openFile(fileName, filterString);
	m_pcap.m_readBufferSize = m_bufferPropertySet.m_readBufferSizeProp.m_value;
	m_captureFilterString = filterString;
	m_state = State.Capturing;

	m_logWriter.write(PcapLogRecordCode.CaptureStarted);
	waitPcapEvents();
	m_stdSessionInfoSet.startSession(m_pluginHost.m_mainThreadScheduler);
	return true;

catch:
	m_logWriter.writeLastError(PcapLogRecordCode.CaptureError);
	return false;
}

void PcapSnifferSession.stopCapture() {
	if (!m_state)
		return;

	m_logWriter.write(PcapLogRecordCode.CaptureStopped);
	m_pcap.close();
	m_state = State.Closed;
	m_syncId++;
	m_stdSessionInfoSet.endSession();
}

bool errorcode PcapSnifferSession.applyCaptureFilter() {
	char const* filter = m_captureFilterProp.m_value;
	if (strcmp(filter, m_captureFilterString) == 0)
		return true;

	m_logWriter.write(PcapLogRecordCode.ChangingCaptureFilter, filter);
	bool result = try m_pcap.setFilter(filter);
	if (!result) {
		m_logWriter.writeLastError(PcapLogRecordCode.CaptureFilterError);
		return false;
	}

	m_logWriter.write(PcapLogRecordCode.CaptureFilterChanged);
	m_captureFilterString = filter;
	return true;
}

bool errorcode PcapSnifferSession.applyViewFilter(bool rebuildLog) {
	m_pluginHost.m_log.suspend();
	bool result = try m_filter.setFilter(m_viewFilterCombo.m_editText);
	m_pluginHost.m_log.resume();

	if (!result) {
		m_viewFilterCombo.m_backgroundColor = ui.StdColor.PastelRed;
		return false;
	}

	m_viewFilterCombo.addEditTextToHistory();
	m_viewFilterCombo.m_backgroundColor = ui.ColorFlags.Undefined;

	if (rebuildLog)
		m_pluginHost.m_log.rebuild();

	return true;
}

void PcapSnifferSession.waitPcapEvents() {
	io.PcapEvents eventMask =
		io.PcapEvents.IoError |
		io.PcapEvents.ReadBufferFull |
		io.PcapEvents.IncomingData |
		io.PcapEvents.Eof;

	m_pcap.wait(eventMask, onPcapEvent~(++m_syncId) @ m_pluginHost.m_mainThreadScheduler);
}

void PcapSnifferSession.load(doc.Storage* storage) {
	m_viewFilterCombo.loadHistory(storage, "viewFilterHistory");
	m_viewFilterCombo.m_editText = storage.readString("viewFilter");
	m_deviceCombo.m_currentIndex = storage.readInt("deviceIdx");
	m_captureFilterProp.m_value = storage.readString("captureFilter");
	m_promisciousProp.m_value = storage.readBool("promiscious", Defaults.Promiscious);
	m_ipFragmentLimitProp.m_value = storage.readInt("ipFragmentLimit", io.IpDefragmenter.Defaults.FragmentLimit);
	m_ipFragmentTimeoutProp.m_value = storage.readInt("ipFragmentTimeout", io.IpDefragmenter.Defaults.FragmentTimeout);
	m_snapshotSizeProp.m_value = storage.readInt("snapshotSize", Defaults.SnapshotSize);
	m_kernelBufferSizeProp.m_value = storage.readInt("kernelBufferSize", Defaults.KernelBufferSise);

	m_bufferPropertySet.load(storage);
	m_readBuffer.setSize(m_bufferPropertySet.m_readBufferSizeProp.m_value);
}

void PcapSnifferSession.save(doc.Storage* storage) {
	m_bufferPropertySet.save(storage);
	m_viewFilterCombo.saveHistory(storage, "viewFilterHistory");
	storage.writeString("viewFilter", m_viewFilterCombo.m_editText);
	storage.writeInt("deviceIdx", (int) m_deviceCombo.m_currentIndex);
	storage.writeString("captureFilter", m_captureFilterProp.m_value);
	storage.writeBool("promiscious", m_promisciousProp.m_value);
	storage.writeInt("ipFragmentLimit", m_ipFragmentLimitProp.m_value);
	storage.writeInt("ipFragmentTimeout", m_ipFragmentTimeoutProp.m_value);
	storage.writeInt("snapshotSize", m_snapshotSizeProp.m_value);
	storage.writeInt("kernelBufferSize", m_kernelBufferSizeProp.m_value);
}

void PcapSnifferSession.updateProperties() {
	m_deviceProp.m_currentIndex = m_deviceCombo.m_currentIndex;
	m_viewFilterProp.m_value = m_viewFilterCombo.m_editText;
}

bool errorcode PcapSnifferSession.applyProperties() {
	m_deviceCombo.m_currentIndex = m_deviceProp.m_currentIndex;
	m_viewFilterCombo.m_editText = m_viewFilterProp.m_value;
	m_ipDefragmenter.m_fragmentLimit = m_ipFragmentLimitProp.m_value;
	m_ipDefragmenter.m_fragmentTimeout = m_ipFragmentTimeoutProp.m_value;
	m_readBuffer.setSize(m_bufferPropertySet.m_readBufferSizeProp.m_value);

	applyViewFilter();

	if (!m_state)
		return true;

	m_pcap.m_readBufferSize = m_bufferPropertySet.m_readBufferSizeProp.m_value;
	applyCaptureFilter();
	return true;
}

void PcapSnifferSession.restoreDefaultProperties() {
	m_promisciousProp.m_value = Defaults.Promiscious;
	m_ipFragmentLimitProp.m_value = io.IpDefragmenter.Defaults.FragmentLimit;
	m_ipFragmentTimeoutProp.m_value = io.IpDefragmenter.Defaults.FragmentTimeout;
	m_snapshotSizeProp.m_value = Defaults.SnapshotSize;
	m_kernelBufferSizeProp.m_value = Defaults.KernelBufferSise;
	m_bufferPropertySet.restoreDefaults();
}

size_t errorcode PcapSnifferSession.transmit(
	void const* p,
	size_t size
) {
	if (!m_pcap.m_isOpen) {
		std.setError("no capture session");
		return -1;
	}

	size_t result = m_pcap.write(p, size);

	io.PcapLinkType linkType = m_pcap.m_linkType;
	void const* params = memcat(&linkType, sizeof(linkType), p, size);
	m_logWriter.write(NetSnifferLogRecordCode.InjectedPacket, params, result + sizeof(linkType));
	m_pluginHost.m_txRxStats.addTx(size);

	return result;
}

void PcapSnifferSession.onPcapEvent(
	uint_t syncId,
	io.PcapEvents triggeredEvents
) {
	if (syncId != m_syncId)
		return;

	if (triggeredEvents & io.PcapEvents.IoError) {
		m_logWriter.write(log.StdRecordCode.Error, m_pcap.m_ioError);
		stopCapture();
		return;
	}

	if ((triggeredEvents & io.PcapEvents.ReadBufferFull) &&
		m_bufferPropertySet.m_notifyReadBufferFullProp.m_value)
		m_logWriter.write(log.StdRecordCode.RxBufferFull);

	if (triggeredEvents & io.PcapEvents.IncomingData) {
		uint64_t timestamp;
		size_t result = m_pcap.read(m_readBuffer.m_p, m_readBuffer.m_size, &timestamp);
		if (result)
			processPacket(timestamp, m_readBuffer.m_p, result);
	}

	if (triggeredEvents & io.PcapEvents.Eof) {
		for (;;) {
			uint64_t timestamp;
			size_t result = m_pcap.read(m_readBuffer.m_p, m_readBuffer.m_size, &timestamp);
			if (!result)
				break;

			processPacket(timestamp, m_readBuffer.m_p, result);
		}

		stopCapture();
		return;
	}

	waitPcapEvents();
}

void PcapSnifferSession.processPacket(
	uint64_t timestamp,
	void const* p,
	size_t size
) {
	io.PcapLinkType linkType = m_pcap.m_linkType;
	void const* params = memcat(&linkType, sizeof(linkType), p, size);
	m_logWriter.write(timestamp, NetSnifferLogRecordCode.Packet, params, size + sizeof(linkType));

	m_pluginHost.m_txRxStats.addRx(size);

	// defragment IPv4

	switch (linkType) {
	case io.PcapLinkType.Loopback:
		if (size < sizeof(io.LoopbackHdr))
			return;

		io.LoopbackHdr const* loopbackHdr = (io.LoopbackHdr const*)p;
		if (loopbackHdr.m_protocolFamily != io.LoopbackProtocolFamily.Ip)
			return;

		p += sizeof(io.LoopbackHdr);
		size -= sizeof(io.LoopbackHdr);
		break;

	case io.PcapLinkType.Ethernet:
		if (size < sizeof(io.EthernetHdr))
			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);
		break;
	}

	if (size < sizeof(io.IpHdr))
		return;

	io.IpHdr const* ipHdr = (io.IpHdr const*) p;
	if ((ipHdr.m_flags & io.IpFlags.Mf) || ipHdr.m_fragmentOffset) {
		ipHdr = m_ipDefragmenter.defragment(timestamp, ipHdr, size);
		if (ipHdr) {
			size = dynamic sizeof(ipHdr);
			m_logWriter.write(NetSnifferLogRecordCode.DefragmentedIp, ipHdr, size);
		}
	}
}

reactor PcapSnifferSession.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";
}

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