//..............................................................................
//
// Serial Monitor over SSH session
//

import "doc_Session.jnc"
import "io_base.jncx"
import "io_ssh.jncx"
import "io_File.jnc"
import "io_Socket.jnc"
import "io_HostNameResolver.jnc"
import "ui_Dialog.jnc"
import "ui_History.jnc"
import "ui_SocketUi.jnc"
import "ui_SshPropertySet.jnc"
import "ui_BufferPropertySet.jnc"
import "ui_StdSessionInfoSet.jnc"
import "log_ThroughputCalc.jnc"
import "log_ChecksumCalc.jnc"
import "SerialMonitor_mi.jnc"
import "std_Buffer.jnc"

using extension ui.ComboBoxHistory;

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

class SshSerialMonSession: doc.Session
{
protected:
	enum Defaults
	{
		ServerPort = 22,
	}

	enum State
	{
		Closed,
		Resolving,
		TcpConnecting,
		SshHandshaking,
		SshAuthenticating,
		WaitingLoginDlg,
		SshOpeningChannel,
		SshRequestingPty,
		SshStartingProcess,
		Connected,
	}

	static char const* m_stateStringTable[] =
	{
		"Closed",
		"Resolving...",
		"Connecting...",
		"Handhsaking...",
		"Authenticating...",
		"Authenticating...",
		"Opening channel...",
		"Requesting PTY...",
		"Starting process...",
		"Connected",
	}

	enum ActionId
	{
		Connect,
		_Count,
	}

	enum IconId
	{
		Connect,
		Disconnect,
		_Count,
	}

	enum StatusPaneId
	{
		RemoteAddress,
		State,
		_Count,
	}

protected:
	// SSH channel

	State bindable m_state;
	io.SshChannel m_ssh;
	io.HostNameResolver m_resolver;
	io.SocketAddress m_remoteAddress;
	std.Buffer m_readBuffer;
	uint_t m_syncId;

	SerialMonitor_mi m_monitor; // device monitor machine interface proxy

	// 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_addressCombo;
	ui.ComboBox* m_portCombo;
	ui.LoginDlg* m_loginDlg;

	// property grid

	ui.SocketPropertySet m_socketPropertySet;
	ui.SshPropertySet m_sshPropertySet;
	ui.BufferPropertySet m_bufferPropertySet;

	// information grid

	ui.StdSessionInfoSet m_stdSessionInfoSet;
	ui.ThroughputInfoSet m_throughputInfoSet;
	ui.ChecksumInfoSet m_checksumInfoSet;

	// ui reactor

	reactor m_uiReactor;

public:
	construct(doc.PluginHost* pluginHost);

	bool errorcode connect();
	bool errorcode connect(io.SocketAddress const* address);
	void disconnect();

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

protected:
	void onConnect()
	{
		m_state ? disconnect() : try connect();
	}

	void onComboEnter()
	{
		try connect();
	}

	void onLoginDlgCompleted(ui.StdDlgButton button);
	void onConsoleSizeChanged();

	void onResolverEvent(
		uint_t syncId,
		io.HostNameResolverEvents triggeredEvents
		);

	void onSshEvent(
		uint_t syncId,
		io.SshEvents triggeredEvents
		);

	void createUi();
	void waitSshEvents();
	void handleConnectError(std.Error const* error);
	void handleSshAuthenticateError(std.Error const* error);
}

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

SshSerialMonSession.construct(doc.PluginHost* pluginHost)
{
	basetype.construct(pluginHost);
	m_monitor.construct(pluginHost.m_log.m_writer);

	m_pluginHost.m_terminal.m_onConsoleSizeChanged += onConsoleSizeChanged;
	m_readBuffer.setSize(ui.BufferPropertySet.Defaults.ReadBufferSize);

	createUi();

	pluginHost.m_log.addSelectionProcessor(log.ThroughputCalcRangeProcessor.factory~(m_throughputInfoSet));
	pluginHost.m_log.addSelectionProcessor(log.ChecksumCalcRangeProcessor.factory~(m_checksumInfoSet));
}

void SshSerialMonSession.createUi()
{
	// toolbar

	m_pluginHost.m_toolBar.addLabel("Address:");
	m_addressCombo = m_pluginHost.m_toolBar.addComboBox(-1);
	m_addressCombo.m_isEditable = true;
	m_addressCombo.m_onEnter += onComboEnter;
	m_addressCombo.m_placeholderText = "Enter a remote IP or hostname...";

	m_pluginHost.m_toolBar.addLabel("Port:");
	m_portCombo = m_pluginHost.m_toolBar.addComboBox(-1);
	m_portCombo.m_isEditable = true;
	m_portCombo.m_onEnter += onComboEnter;
	m_portCombo.m_placeholderText = "Enter a serial port name...";

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

	m_actionTable[ActionId.Connect] = createAction("Connect && capture", m_iconTable [IconId.Connect]);
	m_actionTable[ActionId.Connect].m_onTriggered += onConnect;

	// status bar

	m_statusPaneTable[StatusPaneId.RemoteAddress] = m_pluginHost.m_statusBar.addPane(, 100);
	m_statusPaneTable[StatusPaneId.RemoteAddress].m_color = ui.StdColor.Blue;
	m_statusPaneTable[StatusPaneId.State] = m_pluginHost.m_statusBar.addPane(, 80);

	// property grid

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

	m_socketPropertySet.createProperties(m_pluginHost.m_propertyGrid, group);
	m_socketPropertySet.m_defaultRemotePort = Defaults.ServerPort;

	group = m_pluginHost.m_propertyGrid.createGroupProperty(
		"SSH connection",
		"SSH-specific connection settings"
		);

	m_sshPropertySet.createProperties(
		m_pluginHost.m_propertyGrid,
		group,,
		ui.SshPropertySet.PropertyId.UserName |
		ui.SshPropertySet.PropertyId.AuthenticationMethod |
		ui.SshPropertySet.PropertyId.PrivateKeyFile |
		ui.SshPropertySet.PropertyId.RememberLastPassword
		);

	m_bufferPropertySet.createPropertiesGrouped(
		m_pluginHost.m_propertyGrid,,,
		ui.BufferPropertySet.PropertyId.All & ~ui.BufferPropertySet.PropertyId.ReadParallelism
		);

	// information grid

	ui.InformationGroup* infoGroup = m_pluginHost.m_informationGrid.createGroup("SSH session");

	m_stdSessionInfoSet.create(m_pluginHost.m_txRxStats, m_pluginHost.m_informationGrid, infoGroup);
	m_throughputInfoSet.createInformationValuesGrouped(m_pluginHost.m_informationGrid);
	m_checksumInfoSet.createInformationValuesGrouped(m_pluginHost.m_informationGrid);

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

	m_uiReactor.start();
	m_sshPropertySet.m_uiReactor.start();
}

bool errorcode SshSerialMonSession.connect()
{
	disconnect();
	startSession();

	char const* addressString = m_addressCombo.m_editText;
	m_addressCombo.addEditTextToHistory();

	io.SocketAddress address;
	bool isDirectAddress = try address.parse(addressString);
	if (isDirectAddress)
	{
		if (!address.m_port)
			address.m_port = m_socketPropertySet.m_remotePortProp.m_value;

		return try connect(&address);
	}

	m_state = State.Resolving;

	m_logWriter.write(
		SocketLogRecordCode.Resolving,
		addressString,
		strlen(addressString)
		);

	io.AddressFamily addressFamily = io.AddressFamily.Undefined;
	if (m_socketPropertySet.m_useLocalAddressProp.m_value)
	{
		io.SocketAddress const* adapterAddress = m_socketPropertySet.m_adapterProp.m_value;
		if (adapterAddress)
			addressFamily = adapterAddress.m_family;
	}

	m_resolver.resolve(addressString, addressFamily);

	m_resolver.wait(
		io.HostNameResolverEvents.Resolved | io.HostNameResolverEvents.Error,
		onResolverEvent~(++m_syncId) @ m_pluginHost.m_mainThreadScheduler
		);

	return true;

catch:
	handleConnectError(std.getLastError());
	return false;
}

bool errorcode SshSerialMonSession.connect(io.SocketAddress const* remoteAddress)
{
	disconnect();
	m_monitor.reset();

	m_ssh.m_readBlockSize = m_bufferPropertySet.m_readBlockSizeProp.m_value;
	m_ssh.m_readBufferSize = m_bufferPropertySet.m_readBufferSizeProp.m_value;
	m_ssh.m_writeBufferSize = m_bufferPropertySet.m_writeBufferSizeProp.m_value;
	m_ssh.m_options = ui.getSocketOptions(m_socketPropertySet, m_bufferPropertySet);

	if (!m_socketPropertySet.m_useLocalAddressProp.m_value)
	{
		m_ssh.open(remoteAddress.m_family); // any address
	}
	else
	{
		io.SocketAddress localAddress;
		io.SocketAddress const* adapterAddress = m_socketPropertySet.m_adapterProp.m_value;
		if (adapterAddress)
			localAddress = *adapterAddress;
		else
			localAddress.m_family = remoteAddress.m_family;

		localAddress.m_port = (uint16_t)m_socketPropertySet.m_localPortProp.m_value;

		m_ssh.open(&localAddress);
	}

	m_state = State.TcpConnecting;
	m_remoteAddress = *remoteAddress;

	m_logWriter.write(
		SocketLogRecordCode.Connecting,
		remoteAddress,
		sizeof(io.SocketAddress)
		);

	void const* privateKey;
	char const* password;

	if (m_sshPropertySet.m_authenticationMethodProp.m_value == ui.AuthenticationMethod.PrivateKey)
	{
		privateKey = loadFile(m_sshPropertySet.m_privateKeyFileProp.m_value);
		password = m_loginDlg && m_sshPropertySet.m_rememberLastPasswordProp.m_value ? m_loginDlg.m_privateKeyPassphrase : null;
	}
	else
	{
		privateKey = null;
		password = m_loginDlg && m_sshPropertySet.m_rememberLastPasswordProp.m_value ? m_loginDlg.m_password : null;
	}

	io.SshConnectParams params;
	params.m_address = *remoteAddress;
	params.m_userName = m_sshPropertySet.m_userNameProp.m_value;
	params.m_privateKey = privateKey;
	params.m_privateKeySize = dynamic sizeof(privateKey);
	params.m_password = password;
	params.m_channelType = "session";
	params.m_processType = "exec";

	// we only use IOCTL descriptors on POSIX, but passing those on windows will not hurt (will simply be ignored)

	params.m_processExtra =
		$"tdevmon -m $(m_portCombo.m_editText) --mi "
		$"--ioctl $(psx.TCSETA):$(sizeof (psx.termio)) "
		$"--ioctl $(psx.TCSETAW):$(sizeof (psx.termio)) "
		$"--ioctl $(psx.TCSETAF):$(sizeof (psx.termio)) "
		$"--ioctl $(psx.TCSETS):$(sizeof (psx.termios)) "
		$"--ioctl $(psx.TCSETSW):$(sizeof (psx.termios)) "
		$"--ioctl $(psx.TCSETSF):$(sizeof (psx.termios)) "
		$"--ioctl $(psx.TCSETS2):$(sizeof (psx.termios2)) "
		$"--ioctl $(psx.TCSETSW2):$(sizeof (psx.termios2)) "
		$"--ioctl $(psx.TCSETSF2):$(sizeof (psx.termios2)) "
		$"--ioctl $(psx.TIOCMGET):$(sizeof (uint_t)) "
		$"--ioctl $(psx.TIOCMBIS):$(sizeof (uint_t)) "
		$"--ioctl $(psx.TIOCMBIC):$(sizeof (uint_t))";

	params.m_processExtraSize = strlen(params.m_processExtra);

	m_ssh.connect(&params);
	waitSshEvents();
	return true;

catch:
	handleConnectError(std.getLastError());
	return false;
}

void SshSerialMonSession.waitSshEvents()
{
	io.SshEvents eventMask = io.SshEvents.IoError;

	switch (m_state)
	{
	case State.TcpConnecting:
		eventMask |= io.SshEvents.TcpConnected;
		break;

	case State.SshHandshaking:
		eventMask |= io.SshEvents.SshHandshakeCompleted;
		break;

	case State.SshAuthenticating:
		eventMask |= io.SshEvents.SshAuthenticateCompleted | io.SshEvents.SshAuthenticateError;
		break;

	case State.SshOpeningChannel:
		eventMask |= io.SshEvents.SshChannelOpened;
		break;

	case State.SshRequestingPty:
		eventMask |= io.SshEvents.SshPtyRequestCompleted;
		break;

	case State.SshStartingProcess:
		eventMask |= io.SshEvents.SshConnectCompleted;
		break;

	case State.Connected:
		eventMask |=
			io.SshEvents.TcpDisconnected |
			io.SshEvents.IncomingData;

		break;
	}

	m_ssh.wait(eventMask, onSshEvent~(++m_syncId) @ m_pluginHost.m_mainThreadScheduler);
}

void SshSerialMonSession.handleConnectError(std.Error const* error)
{
	m_logWriter.write(SocketLogRecordCode.ConnectError, error);
	m_ssh.close();
	m_state = State.Closed;
}

void SshSerialMonSession.handleSshAuthenticateError(std.Error const* error)
{
	m_logWriter.write(SocketLogRecordCode.ConnectError, error);

	if (!m_loginDlg)
	{
		m_loginDlg = new ui.LoginDlg;
		m_loginDlg.m_onCompleted += onLoginDlgCompleted;
	}

	m_loginDlg.m_userName = m_sshPropertySet.m_userNameProp.m_value;
	m_loginDlg.m_authenticationMethod = m_sshPropertySet.m_authenticationMethodProp.m_value;
	m_loginDlg.m_privateKeyFileName = m_sshPropertySet.m_privateKeyFileProp.m_value;

	if (!m_sshPropertySet.m_rememberLastPasswordProp.m_value)
	{
		m_loginDlg.m_password = null;
		m_loginDlg.m_privateKeyPassphrase = null;
	}

	m_loginDlg.show();

	m_state = State.WaitingLoginDlg;
}

void SshSerialMonSession.onLoginDlgCompleted(ui.StdDlgButton button)
{
	if (m_state != State.WaitingLoginDlg)
		return;

	if (button != ui.StdDlgButton.Ok)
	{
		m_ssh.close();
		m_state = State.Closed;
		return;
	}

	State nextState;

	if (strcmp(m_loginDlg.m_userName, m_sshPropertySet.m_userNameProp.m_value) == 0)
	{
		nextState = State.SshAuthenticating;
		m_logWriter.write(SshLogRecordCode.SshAuthenticating);
	}
	else
	{
		nextState = State.TcpConnecting;
		m_logWriter.write(SshLogRecordCode.Reconnecting);
	}

	switch (m_loginDlg.m_authenticationMethod)
	{
	case ui.AuthenticationMethod.Password:
		m_ssh.authenticate(m_loginDlg.m_userName, null, 0, m_loginDlg.m_password);
		break;

	case ui.AuthenticationMethod.PrivateKey:
		void const* privateKey = loadFile(m_loginDlg.m_privateKeyFileName);

		m_ssh.authenticate(
			m_loginDlg.m_userName,
			privateKey,
			dynamic sizeof(privateKey),
			m_loginDlg.m_privateKeyPassphrase
			);
		break;
	}

	m_sshPropertySet.m_userNameProp.m_value = m_loginDlg.m_userName;
	m_sshPropertySet.m_authenticationMethodProp.m_value = m_loginDlg.m_authenticationMethod;
	m_sshPropertySet.m_privateKeyFileProp.m_value = m_loginDlg.m_privateKeyFileName;
	m_state = nextState;
	waitSshEvents();

catch:
	m_ssh.close();
	m_state = State.Closed;
}

void SshSerialMonSession.onConsoleSizeChanged()
{
	m_ssh.resizePty(
		m_pluginHost.m_terminal.m_consoleWidth,
		m_pluginHost.m_terminal.m_consoleHeight
		);
}

void SshSerialMonSession.disconnect()
{
	switch (m_state)
	{
	case State.Closed:
		return;

	case State.Resolving:
	case State.TcpConnecting:
	case State.SshHandshaking:
	case State.SshAuthenticating:
	case State.WaitingLoginDlg:
	case State.SshOpeningChannel:
	case State.SshRequestingPty:
	case State.SshStartingProcess:
		m_logWriter.write(SocketLogRecordCode.ConnectCancelled);
		break;

	case State.Connected:
		TcpDisconnectLogRecordParams params;
		params.m_address = m_remoteAddress;

		m_logWriter.write(
			SocketLogRecordCode.Disconnected,
			&params,
			sizeof(params)
			);
		break;
	}

	m_ssh.close();
	m_state = State.Closed;
	m_syncId++;
	m_stdSessionInfoSet.endSession();
}

void SshSerialMonSession.load(doc.Storage* storage)
{
	m_addressCombo.loadHistory(storage, "addressHistory");
	m_addressCombo.m_editText = storage.readString("address");

	m_portCombo.loadHistory(storage, "portHistory");
	m_portCombo.m_editText = storage.readString("port");

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

void SshSerialMonSession.save(doc.Storage* storage)
{
	m_socketPropertySet.save(storage);
	m_sshPropertySet.save(storage);
	m_bufferPropertySet.save(storage);

	m_addressCombo.saveHistory(storage, "addressHistory");
	storage.writeString("address", m_addressCombo.m_editText);

	m_portCombo.saveHistory(storage, "portHistory");
	storage.writeString("port", m_portCombo.m_editText);
}

void SshSerialMonSession.updateProperties()
{
	m_socketPropertySet.m_remoteAddressProp.m_value = m_addressCombo.m_editText;
}

bool errorcode SshSerialMonSession.applyProperties()
{
	m_addressCombo.m_editText = m_socketPropertySet.m_remoteAddressProp.m_value;
	m_readBuffer.setSize(m_bufferPropertySet.m_readBufferSizeProp.m_value);

	if (!m_ssh.m_isOpen)
		return true;

	m_ssh.m_options = ui.getSocketOptions(m_socketPropertySet, m_bufferPropertySet);
	return true;
}

void SshSerialMonSession.restoreDefaultProperties()
{
	m_socketPropertySet.restoreDefaults();
	m_sshPropertySet.restoreDefaults();
	m_bufferPropertySet.restoreDefaults();
}

void SshSerialMonSession.onSshEvent(
	uint_t syncId,
	io.SshEvents triggeredEvents
	)
{
	if (syncId != m_syncId)
		return;

	if (triggeredEvents & io.SshEvents.IoError)
	{
		handleConnectError(m_ssh.m_ioError);
		return;
	}

	if (triggeredEvents & io.SshEvents.SshAuthenticateError)
	{
		handleSshAuthenticateError(m_ssh.m_ioError);
		return;
	}

	if (triggeredEvents & io.SshEvents.TcpConnected)
	{
		m_logWriter.write(SshLogRecordCode.TcpConnectCompleted);
		m_state = State.SshHandshaking;
	}

	if (triggeredEvents & io.SshEvents.SshHandshakeCompleted)
	{
		m_logWriter.write(SshLogRecordCode.SshHandshakeCompleted);
		m_state = State.SshAuthenticating;
	}

	if (triggeredEvents & io.SshEvents.SshAuthenticateCompleted)
	{
		m_logWriter.write(SshLogRecordCode.SshAuthCompleted);
		m_state = State.SshOpeningChannel;
	}

	if (triggeredEvents & io.SshEvents.SshChannelOpened)
	{
		m_logWriter.write(SshLogRecordCode.SshChannelOpened);
		m_state = State.SshRequestingPty;
	}

	if (triggeredEvents & io.SshEvents.SshPtyRequestCompleted)
	{
		m_logWriter.write(SshLogRecordCode.SshPtyRequested);
		m_state = State.SshStartingProcess;
	}

	if (triggeredEvents & io.SshEvents.SshConnectCompleted)
	{
		m_logWriter.write(SshLogRecordCode.SshProcessStarted);
		m_logWriter.write(SshLogRecordCode.ConnectCompleted);
		m_state = State.Connected;
		m_stdSessionInfoSet.startSession(m_pluginHost.m_mainThreadScheduler);
	}

	if (triggeredEvents & io.SshEvents.TcpDisconnected)
	{
		for (;;) // read all the buffered data
		{
			size_t size = m_ssh.read(m_readBuffer.m_p, m_readBuffer.m_size);
			if (!size)
				break;

			m_monitor.write(m_readBuffer.m_p, size);
		}

		if (!m_monitor.m_msgCount)
		{
			if (m_monitor.m_buffer.m_size)
				m_logWriter.write(log.StdRecordCode.Bin, m_monitor.m_buffer.m_p, m_monitor.m_buffer.m_size);

			std.setError("invalid or missing tdevmon on the target machine");
			m_logWriter.writeLastError(log.StdRecordCode.Error);
		}

		TcpDisconnectLogRecordParams params;
		params.m_address = m_remoteAddress;
		params.m_flags = TcpDisconnectLogRecordFlags.Remote;

		if (triggeredEvents & io.SshEvents.TcpReset)
			params.m_flags |= TcpDisconnectLogRecordFlags.Reset;

		m_logWriter.write(
			SocketLogRecordCode.Disconnected,
			&params,
			sizeof(params)
			);

		m_ssh.close();
		m_state = State.Closed;
		m_stdSessionInfoSet.endSession();
		return;
	}

	if (triggeredEvents & io.SshEvents.IncomingData)
	{
		size_t size = m_ssh.read(m_readBuffer.m_p, m_readBuffer.m_size);
		if (size)
			m_monitor.write(m_readBuffer.m_p, size);
	}

	waitSshEvents();

catch:
	m_logWriter.writeLastError(log.StdRecordCode.Error);
	disconnect();
}

void SshSerialMonSession.onResolverEvent(
	uint_t syncId,
	io.HostNameResolverEvents triggeredEvents
	)
{
	if (syncId != m_syncId)
		return;

	if (triggeredEvents & io.HostNameResolverEvents.Error)
	{
		handleConnectError(m_resolver.m_error);
		return;
	}

	m_logWriter.write(
		SocketLogRecordCode.ResolveCompleted,
		m_resolver.m_addressTable,
		sizeof(io.SocketAddress)
		);

	m_state = State.Closed;

	io.SocketAddress address = m_resolver.m_addressTable[0];
	if (!address.m_port)
		address.m_port = m_socketPropertySet.m_remotePortProp.m_value;

	try connect(address);
}

reactor SshSerialMonSession.m_uiReactor
{
	m_title = $"Mon $(m_addressCombo.m_currentText)";
	m_actionTable[ActionId.Connect].m_text = m_state ? "Disconnect" : "Connect && capture";
	m_actionTable[ActionId.Connect].m_icon = m_iconTable[m_state ? IconId.Disconnect : IconId.Connect];
	m_statusPaneTable[StatusPaneId.State].m_text = m_stateStringTable[m_state];
	m_statusPaneTable[StatusPaneId.RemoteAddress].m_text = m_state > State.Resolving ? m_remoteAddress.getString() : "<peer-address>";
	m_statusPaneTable[StatusPaneId.RemoteAddress].m_isVisible = m_state > State.Resolving;
}

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

// simple helper to load private key files (returning null is OK)

void const* loadFile(char const* fileName)
{
	if (!fileName || !*fileName)
		return null;

	io.File file;
	file.open(fileName, io.FileOpenFlags.ReadOnly);
	size_t size = file.m_size; // cache size (file may change)
	void* p = new char[size];
	file.read(p, size);
	return p;

catch:
	return null;
}

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