//..............................................................................
//
// SSH Connection session
//

import "io_base.jncx"
import "io_ssh.jncx"
import "io_Socket.jnc"
import "io_MappedFile.jnc"
import "io_SocketAddressResolver.jnc"
import "std_Buffer.jnc"
import "doc_Plugin.jnc"
import "gui_Dialog.jnc"
import "gui_History.jnc"

using extension gui.ComboBoxHistory;

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

class SshChannelSession: doc.Session
{
protected:
	enum Defaults
	{
		ServerPort           = 22,
		AuthenticationMethod = gui.AuthenticationMethod.Password,
		RememberLastPassword = false,
	}

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

	static char DefaultChannelType [] = "session";
	static char DefaultProcessType [] = "shell";
	static char DefaultPtyType [] = "xterm";

	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.SocketAddressResolver m_resolver;
	io.SocketAddress m_remoteAddress;
	std.Buffer m_readBuffer;
	uint_t m_syncId;

	// 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_addressCombo;
	gui.LoginDlg* m_loginDlg;

	// property grid

	gui.StringProperty* m_userNameProp;
	gui.EnumProperty* m_authenticationMethodProp;
	gui.FileProperty* m_privateKeyFileProp;
	gui.BoolProperty* m_rememberLastPasswordProp;
	gui.StringProperty* m_channelTypeProp;
	gui.StringProperty* m_processTypeProp;
	gui.StringProperty* m_ptyTypeProp;
	io.SocketPropertySet m_socketPropertySet;
	io.BufferPropertySet m_bufferPropertySet;

public:
	construct (doc.PluginHost* pluginHost);

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

	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:
	onConnect ()
	{
		m_state ? disconnect () : try connect ();
	}

	onAddressComboEnter ()
	{
		try connect ();
	}

	onResolveCompleted (
		uint_t syncId,
		io.SocketAddress const* addressTable,
		size_t addressCount,
		std.Error const* error
		);

	onLoginDlgCompleted (gui.StdDlgButton button);
	onConsoleSizeChanged ();

	onSshEvent (
		uint_t syncId,
		io.SshEvents triggeredEvents
		);

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

	reactor m_uiReactor ();
}

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

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

	pluginHost.m_log.attachRepresenter (representSshChannelLog);

	m_state = State.Closed;
	m_pluginHost.m_terminal.m_onConsoleSizeChanged += onConsoleSizeChanged;
	m_readBuffer.setSize (io.BufferPropertySet.Defaults.ReadBufferSize);

	createUi ();
}

SshChannelSession.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 += onAddressComboEnter;

	m_iconTable [IconId.Connect] = new gui.Icon ("images/connect.png");
	m_iconTable [IconId.Disconnect] = new gui.Icon ("images/disconnect.png");

	m_actionTable [ActionId.Connect] = createAction ("Connect", 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 = gui.StdColor.Blue;
	m_statusPaneTable [StatusPaneId.State] = m_pluginHost.m_statusBar.addPane (, 80);

	// property grid

	gui.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_userNameProp = m_pluginHost.m_propertyGrid.createStringProperty (
		group,,
		"User name",
		"User name for SSH authentification"
		);

	static gui.EnumPropertyOption authenticationMethodTable [] =
	{
		{ "Password",    gui.AuthenticationMethod.Password },
		{ "Private key", gui.AuthenticationMethod.PrivateKey },
	}

	m_authenticationMethodProp = m_pluginHost.m_propertyGrid.createEnumProperty (
		group,,
		"User authentication",
		"How to authentify a user (password / private key)",
		authenticationMethodTable,
		countof (authenticationMethodTable)
		);

	m_privateKeyFileProp = m_pluginHost.m_propertyGrid.createFileProperty (
		group,,
		"Private key file",
		"Private key file (PEM) for SSH authentification"
		);

	m_privateKeyFileProp.m_fileDialogCaption = "Select Private Key File";
	m_privateKeyFileProp.m_filter = "SSH identity files (id_dsa id_rsa);;PEM files (*.pem);;All files (*.*)";

	m_rememberLastPasswordProp = m_pluginHost.m_propertyGrid.createBoolProperty (
		group,,
		"Remember last password",
		"Re-use the last entered password during subsequent connect attempts"
		);

	m_channelTypeProp = m_pluginHost.m_propertyGrid.createStringProperty (
		group,,
		"SSH channel",
		"Type of requested SSH channel"
		);

	m_processTypeProp = m_pluginHost.m_propertyGrid.createStringProperty (
		group,,
		"SSH process",
		"Type of requested SSH process"
		);

	m_ptyTypeProp = m_pluginHost.m_propertyGrid.createStringProperty (
		group,,
		"Pseudoterminal",
		"Type of requested pseudoterminal (PTY)"
		);

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

	m_bufferPropertySet.createProperties (
		m_pluginHost.m_propertyGrid,
		group,,
		io.BufferPropertySet.PropertyId.All & ~io.BufferPropertySet.PropertyId.ReadParallelism
		);

	m_uiReactor.start ();
}

bool errorcode SshChannelSession.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_ip4.m_port) // IPv4 & IPv6 share the same port location
			address.m_ip4.m_port = m_socketPropertySet.m_remotePortProp.m_value;

		return try connect (&address);
	}

	m_state = State.Resolving;

	writeLog (
		SshChannelLogRecordCode.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, onResolveCompleted ~(++m_syncId) @ m_pluginHost.m_mainThreadScheduler);
	return true;

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

bool errorcode SshChannelSession.connect (io.SocketAddress const* remoteAddress)
{
	disconnect ();

	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 = io.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_ip4.m_port = (uint16_t) m_socketPropertySet.m_localPortProp.m_value; // same for IPv4/IPv6

		m_ssh.open (&localAddress);
	}

	m_state = State.TcpConnecting;
	m_remoteAddress = *remoteAddress;

	writeLog (
		SshChannelLogRecordCode.TcpConnecting,
		remoteAddress,
		sizeof (io.SocketAddress)
		);

	void const* privateKey;
	char const* password;

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

	m_ssh.connect (
		remoteAddress,
		m_userNameProp.m_value,
		privateKey,
		dynamic sizeof (privateKey),
		password,
		m_channelTypeProp.m_value,
		m_processTypeProp.m_value,
		m_ptyTypeProp.m_value,
		m_pluginHost.m_terminal.m_consoleWidth,
		m_pluginHost.m_terminal.m_consoleHeight
		);

	waitSshEvents ();
	return true;

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

SshChannelSession.waitSshEvents ()
{
	io.SshEvents eventMask = io.SshEvents.IoError;

	switch (m_state)
	{
	case State.TcpConnecting:
		eventMask |= io.SshEvents.TcpConnectCompleted;
		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;

		if (!m_isTransmitEnabled)
			eventMask |= io.SshEvents.WriteBufferReady;

		break;
	}

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

SshChannelSession.handleConnectError (std.Error const* error)
{
	writeLog (
		SshChannelLogRecordCode.ConnectError,
		error,
		error.m_size
		);

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

SshChannelSession.handleSshAuthenticateError (std.Error const* error)
{
	writeLog (
		SshChannelLogRecordCode.ConnectError,
		error,
		error.m_size
		);

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

	m_loginDlg.m_userName = m_userNameProp.m_value;
	m_loginDlg.m_authenticationMethod = m_authenticationMethodProp.m_value;
	m_loginDlg.m_privateKeyFileName = m_privateKeyFileProp.m_value;
	m_loginDlg.show ();

	m_state = State.WaitingLoginDlg;
}

SshChannelSession.onLoginDlgCompleted (gui.StdDlgButton button)
{
	if (m_state != State.WaitingLoginDlg)
		return;

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

	State nextState;

	if (strcmp (m_loginDlg.m_userName, m_userNameProp.m_value) == 0)
	{
		nextState = State.SshAuthenticating;
		writeLog (SshChannelLogRecordCode.SshAuthenticating);
	}
	else
	{
		nextState = State.TcpConnecting;
		writeLog (SshChannelLogRecordCode.Reconnecting);
	}

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

	case gui.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_userNameProp.m_value = m_loginDlg.m_userName;
	m_authenticationMethodProp.m_value = m_loginDlg.m_authenticationMethod;
	m_privateKeyFileProp.m_value = m_loginDlg.m_privateKeyFileName;
	m_state = nextState;
	waitSshEvents ();

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

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

SshChannelSession.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:
		writeLog (SshChannelLogRecordCode.ConnectCancelled);
		break;

	case State.Connected:
		writeLog (
			SshChannelLogRecordCode.Disconnected,
			&m_remoteAddress,
			sizeof (m_remoteAddress)
			);
		break;
	}

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

void SshChannelSession.load (doc.Storage* storage)
{
	m_addressCombo.loadHistory (storage, "addressHistory");
	m_addressCombo.m_editText = storage.readString ("address");
	m_userNameProp.m_value = storage.readString ("userName");
	m_authenticationMethodProp.m_value = storage.readInt ("authenticationMethod", Defaults.AuthenticationMethod);
	m_privateKeyFileProp.m_value = storage.readString ("privateKeyFile");
	m_rememberLastPasswordProp.m_value = storage.readBool ("rememberLastPassword", Defaults.RememberLastPassword);
	m_channelTypeProp.m_value = storage.readString ("channelType", DefaultChannelType);
	m_processTypeProp.m_value = storage.readString ("processType", DefaultProcessType);
	m_ptyTypeProp.m_value = storage.readString ("ptyType", DefaultPtyType);

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

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

	m_addressCombo.saveHistory (storage, "addressHistory");
	storage.writeString ("address", m_addressCombo.m_editText);
	storage.writeString ("userName", m_userNameProp.m_value);
	storage.writeInt ("authenticationMethod", m_authenticationMethodProp.m_value);
	storage.writeString ("privateKeyFile", m_privateKeyFileProp.m_value);
	storage.writeBool ("rememberLastPassword", m_rememberLastPasswordProp.m_value);
	storage.writeString ("channelType", m_channelTypeProp.m_value);
	storage.writeString ("processType", m_processTypeProp.m_value);
	storage.writeString ("ptyType", m_ptyTypeProp.m_value);
}

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

bool errorcode SshChannelSession.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 = io.getSocketOptions (m_socketPropertySet, m_bufferPropertySet);
	return true;
}

void SshChannelSession.restoreDefaultProperties ()
{
	m_rememberLastPasswordProp.m_value = Defaults.RememberLastPassword;
	m_channelTypeProp.m_value = DefaultChannelType;
	m_processTypeProp.m_value = DefaultProcessType;
	m_ptyTypeProp.m_value = DefaultPtyType;

	m_socketPropertySet.restoreDefaults ();
	m_bufferPropertySet.restoreDefaults ();
}

size_t errorcode SshChannelSession.transmit (
	void const* p,
	size_t size
	)
{
	if (m_state != State.Connected)
		return -1;

	size_t result = m_ssh.write (p, size);
	if (result)
		writeLog (log.StdRecordCode.Tx, p, result);

	if (result < size)
	{
		m_isTransmitEnabled = false;
		waitSshEvents ();
	}

	return result;
}

SshChannelSession.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.TcpConnectCompleted)
	{
		writeLog (SshChannelLogRecordCode.TcpConnectCompleted);
		m_state = State.SshHandshaking;
	}

	if (triggeredEvents & io.SshEvents.SshHandshakeCompleted)
	{
		writeLog (SshChannelLogRecordCode.SshHandshakeCompleted);
		m_state = State.SshAuthenticating;
	}

	if (triggeredEvents & io.SshEvents.SshAuthenticateCompleted)
	{
		writeLog (SshChannelLogRecordCode.SshAuthCompleted);
		m_state = State.SshOpeningChannel;
	}

	if (triggeredEvents & io.SshEvents.SshChannelOpened)
	{
		writeLog (SshChannelLogRecordCode.SshChannelOpened);
		m_state = State.SshRequestingPty;
	}

	if (triggeredEvents & io.SshEvents.SshPtyRequestCompleted)
	{
		writeLog (SshChannelLogRecordCode.SshPtyRequested);
		m_state = State.SshStartingProcess;
	}

	if (triggeredEvents & io.SshEvents.SshConnectCompleted)
	{
		writeLog (SshChannelLogRecordCode.SshProcessStarted);
		writeLog (SshChannelLogRecordCode.ConnectCompleted);
		m_state = State.Connected;
		m_isTransmitEnabled = true;
	}

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

			writeLog (log.StdRecordCode.Rx, m_readBuffer.m_p, result);
		}

		writeLog (
			SshChannelLogRecordCode.Disconnected,
			&m_remoteAddress,
			sizeof (m_remoteAddress)
			);

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

	if (triggeredEvents & io.SshEvents.IncomingData)
	{
		size_t result = m_ssh.read (m_readBuffer.m_p, m_readBuffer.m_size);
		if (result)
			writeLog (log.StdRecordCode.Rx, m_readBuffer.m_p, result);
	}

	if (io.SshEvents.WriteBufferReady)
		m_isTransmitEnabled = true;

	waitSshEvents ();
}

SshChannelSession.onResolveCompleted (
	uint_t syncId,
	io.SocketAddress const* addressTable,
	size_t addressCount,
	std.Error const* error
	)
{
	if (syncId != m_syncId)
		return;

	if (error)
	{
		handleConnectError (error);
		return;
	}
	writeLog (
		SshChannelLogRecordCode.ResolveCompleted,
		addressTable,
		sizeof (io.SocketAddress)
		);

	m_state = State.Closed;

	io.SocketAddress address = addressTable [0];
	if (!address.m_ip4.m_port) // IPv4 & IPv6 share the same port location
		address.m_ip4.m_port = m_socketPropertySet.m_remotePortProp.m_value;

	connect (&address);
}

reactor SshChannelSession.m_uiReactor ()
{
	m_title = $"SSH $(m_addressCombo.m_currentText)";
	m_privateKeyFileProp.m_isEnabled = m_authenticationMethodProp.m_value == gui.AuthenticationMethod.PrivateKey;
	m_socketPropertySet.m_adapterProp.m_isEnabled = m_socketPropertySet.m_useLocalAddressProp.m_value;
	m_socketPropertySet.m_localPortProp.m_isEnabled = m_socketPropertySet.m_useLocalAddressProp.m_value;
	m_socketPropertySet.m_reuseLocalAddressProp.m_isEnabled = m_socketPropertySet.m_useLocalAddressProp.m_value;
	m_actionTable [ActionId.Connect].m_text = m_state ? "Disconnect" : "Connect";
	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;
}

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