From 33884f1bffc2c1783b349aecec36faccafc58590 Mon Sep 17 00:00:00 2001
From: Muthu Kumar <muthukumar@thefeathers.in>
Date: Tue, 17 Jul 2018 14:44:28 +0530
Subject: [PATCH] [tsh] Added /start and /save by identifier

---
 README.md              |  2 --
 app.js                 | 30 +++++++++++++++++++++++-------
 lib/responseHandler.js |  6 +++++-
 lib/sessionFinder.js   | 10 +++-------
 util/index.js          | 11 ++++++++++-
 5 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/README.md b/README.md
index 462e24c..e2e8677 100644
--- a/README.md
+++ b/README.md
@@ -33,8 +33,6 @@ Try running `/start` in your bot's private.
 
 - `tsh` doesn't play well with commands that require password from stdin, like `sudo`, or any command that creates a new shell (like `bash`). As a workaround, you could use `sudo -S` to read password from stdin. These may be fixed in a later version.
 
-- Some times, the same response will be sent several times as the stream is triggered. This is fixable, and will be done soon.
-
 ## Security
 
 `tsh` is recommended to be used by a single user only, for your personal needs. Therefore, you must host your own bot on the server. Support for multiple machines may come in future.
diff --git a/app.js b/app.js
index 4815471..aa82d43 100644
--- a/app.js
+++ b/app.js
@@ -15,6 +15,9 @@ const responder = require('./lib/responseHandler.js');
 const sessionFinder = require('./lib/sessionFinder.js');
 const listeners = require('./lib/listeners.js');
 
+// Utils
+const { extractCommandText } = require('./util/index.js');
+
 const dateOptions = {
 	weekday: 'long',
 	year: 'numeric',
@@ -45,20 +48,33 @@ bot.command('start',
 		const newProc = spawn(defaultShell, {
 			cwd: home
 		});
-		newProc.identifier = identifierState;
+		const identifier = extractCommandText('start')(ctx);
+		if(identifier) newProc.identifier = identifier;
+		else newProc.identifier = identifierState;
+		newProc.index = identifierState;
 		sessions[identifierState] = newProc;
 		identifierState++;
 		sessions.currentSession = newProc;
 		listeners.add(sessions.currentSession, responder, ctx);
-		return ctx.replyWithHTML(`Welcome to tsh -- <code>Telegram Shell!</code>\n\n`
+		return responder.success(`Welcome to tsh -- <code>Telegram Shell!</code>\n\n`
 			+ `You are now connected to <code>${hostname}</code>`
-			+ ` as <strong>${username}</strong>`);
+			+ ` as <strong>${username}</strong>.`,
+			'html'
+		)(ctx);
+	});
+
+bot.command('save',
+	ctx => {
+		const identifier = extractCommandText('save')(ctx);
+		if(!identifier) return responder.fail('Need a valid identifier to save session.')(ctx);
+		sessions.currentSession.identifier = identifier;
+		return responder.success(`Saved session <code>${identifier}</code>.`, 'html')(ctx);
 	});
 
 bot.command('ls',
 	ctx => ctx.reply(
-		sessions.reduce((acc, _, index) =>
-			acc ? `${acc}\n${index}` : `${index}`, '')
+		sessions.reduce((acc, session) =>
+			acc ? `${acc}\n${session.identifier}` : `${session.identifier}`, '')
 		|| `No sessions found. Start one with /start.`
 	));
 
@@ -87,9 +103,9 @@ bot.command('kill',
 		const session = getSession(ctx)('kill') || sessions.currentSession;
 		if(!session)
 			return responder.fail('Session not found. /ls for list of sessions.')(ctx);
-		delete sessions[session.identifier];
-		if(session === sessions.currentSession) sessions.currentSession = undefined;
 		session.kill();
+		delete sessions[session.index];
+		if(session === sessions.currentSession) sessions.currentSession = undefined;
 		ctx.reply('Session killed. /ls for list of sessions.')
 	})
 
diff --git a/lib/responseHandler.js b/lib/responseHandler.js
index eeb4ba1..340e802 100644
--- a/lib/responseHandler.js
+++ b/lib/responseHandler.js
@@ -1,7 +1,11 @@
-const success = response => ctx => response ? ctx.reply(response) : null;
 const { EOL } = require('os');
 const { path } = require('../util/index.js');
 
+const success = (response, mode) => ctx => {
+	const respondAs = (mode && mode.toLowerCase() === 'html') ? 'replyWithHTML' : 'reply';
+	return response ? ctx[respondAs](response) : null;
+};
+
 const convertCtx = ctx => {
 	const message = path(['update', 'message'], ctx);
 	const { id, first_name, username, language_code } = path(['from'], message);
diff --git a/lib/sessionFinder.js b/lib/sessionFinder.js
index a086385..164c7e5 100644
--- a/lib/sessionFinder.js
+++ b/lib/sessionFinder.js
@@ -1,12 +1,8 @@
-const { compose, getText, extractCommandText } = require('../util/index.js');
+const { extractCommandText } = require('../util/index.js');
 
 const sessionFinder = sessions => ctx => command => {
-	const identifier = compose(
-		extractCommandText(command),
-		getText,
-	)(ctx);
-
-	return sessions[identifier];
+	const identifier = extractCommandText(command)(ctx);
+	return sessions[identifier] || sessions.filter(session => session.identifier === identifier)[0];
 };
 
 module.exports = sessionFinder;
\ No newline at end of file
diff --git a/util/index.js b/util/index.js
index 934ce60..f716110 100644
--- a/util/index.js
+++ b/util/index.js
@@ -1,7 +1,16 @@
 const path = (path, obj) => path.reduce((result, segment) => result && result[segment], obj);
+
 const compose = (...fns) => fns.reduce((f, g) => (...args) => f(g(...args)));
+
 const getText = ctx => path(['update', 'message', 'text'], ctx) || '';
-const extractCommandText = cmd => text => text.replace(`/${cmd}`, '').trim();
+
+const removeCommand = cmd => text => text.replace(`/${cmd}`, '').trim();
+
+const extractCommandText = cmd => ctx =>
+	compose(
+		removeCommand(cmd),
+		getText,
+	)(ctx);
 
 module.exports = {
 	path,