Implement password authentication

.gitignore

@@ -128,4 +128,4 @@ dmypy.json
 # Pyre type checker
 .pyre/
 
-token.txt
+config.py

README.md

@@ -2,9 +2,6 @@
 
 This program is a PC kill switch Telegram bot. It is a project for New Generation.
 
-## Security notice
-This program provides no access control, so anyone knowing the bot username is able to use it.
-
 ## Features
 - disabling and re-enabling keyboard and mouse input
 - locking the screen
@@ -17,7 +14,7 @@ This program provides no access control, so anyone knowing the bot username is a
 - works on Windows only
 
 ## Usage
-Before starting the bot, its token must be put into token.txt file with no newlines. Administrator privileges are required so that the bot is able to disable input and format partitions. Once the bot is started, send /start to its account on Telegram to see what can be done.
+Before starting the bot, copy config-example.py file to config.py and put the password and the token. Administrator privileges are required so that the bot is able to disable input and format partitions. Once the bot is started, send /start to its account on Telegram and enter the password to see what can be done.
 
 ## Licensing
 All code in this repository is Unlicensed, see UNLICENSE.

config-example.py

@@ -0,0 +1,11 @@
+# Project: Telegram PC kill switch bot
+# Author: Stanislav Mykhailenko
+# License: Unlicense
+
+# This file contains an example config for the program
+
+# Telegram bot token
+token = "put your token here"
+
+# Bot password
+master_password = "put your password here"

tg.py

@@ -4,6 +4,7 @@
 
 # This file contains Telegram interactions
 
+from config import *
 from telegram import ReplyKeyboardMarkup, ReplyKeyboardRemove, Update
 from telegram.ext import Application, CommandHandler, ContextTypes, ConversationHandler, filters, MessageHandler
 from payload import *
@@ -16,14 +17,14 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-ACTION, FORMAT, MESSAGE = range(3)
-
-def getToken():
-	with open(os.path.realpath(os.path.dirname(__file__)) + '/' + 'token.txt', 'r') as file:
-		return file.read()
-
+PASSWORD, ACTION, FORMAT, MESSAGE = range(4)
 
 async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
+    await update.message.reply_text("Enter password.")
+
+    return PASSWORD
+
+async def show_keyboard(update: Update):
 	reply_keyboard = [["Enable input", "Disable input", "Lock screen", "Format volumes", "Play message"]]
 
 	await update.message.reply_text(
@@ -33,8 +34,20 @@ async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
 		),
 	)
 
+async def password(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
+	user = update.message.from_user
+
+	if update.message.text == master_password:
+		logger.info("User %s logged in.", user.first_name)
+		await show_keyboard(update)
+
 		return ACTION
 
+	logger.info("User %s failed to log in.", user.first_name)
+	await update.message.reply_text("Wrong password entered, please try again.")
+
+	return PASSWORD
+
 async def action(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
 	user = update.message.from_user
 
@@ -43,17 +56,20 @@ async def action(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
 			logger.info("Got a request from %s to enable input.", user.first_name)
 			enableInput()
 			await update.message.reply_text("Input enabled.", reply_markup=ReplyKeyboardRemove())
-			return ConversationHandler.END
+			await show_keyboard(update)
+			return ACTION
 		case "Disable input":
 			logger.info("Got a request from %s to disable input.", user.first_name)
 			disableInput()
 			await update.message.reply_text("Input disabled.", reply_markup=ReplyKeyboardRemove())
-			return ConversationHandler.END
+			await show_keyboard(update)
+			return ACTION
 		case "Lock screen":
 			logger.info("Got a request from %s to lock screen.", user.first_name)
 			lockScreen()
 			await update.message.reply_text("Screen locked.", reply_markup=ReplyKeyboardRemove())
-			return ConversationHandler.END
+			await show_keyboard(update)
+			return ACTION
 		case "Format volumes":
 			logger.info("Got a request from %s to format volumes.", user.first_name)
 			await update.message.reply_text("Please type a space-separated list of the volumes you want to format.", reply_markup=ReplyKeyboardRemove())
@@ -70,7 +86,8 @@ async def format(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
 	Thread(target=formatVolumes,args=(update.message.text.split(),)).start()
 	await update.message.reply_text("Command to format the volumes sent.")
 
-	return ConversationHandler.END
+	await show_keyboard(update)
+	return ACTION
 
 async def message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
 	user = update.message.from_user
@@ -78,21 +95,23 @@ async def message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
 	Thread(target=playMessage,args=(update.message.text,)).start()
 	await update.message.reply_text("Command to play the message sent.")
 
-	return ConversationHandler.END
+	await show_keyboard(update)
+	return ACTION
 
 async def cancel(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
 	user = update.message.from_user
-	logger.info("User %s canceled the conversation.", user.first_name)
+	logger.info("User %s cancelled the conversation.", user.first_name)
 	await update.message.reply_text("Request cancelled.", reply_markup=ReplyKeyboardRemove())
 
 	return ConversationHandler.END
 
 def startBot():
-	application = Application.builder().token(getToken()).build()
+	application = Application.builder().token(token).build()
 	
 	conv_handler = ConversationHandler(
 		entry_points=[CommandHandler("start", start)],
 		states={
+			PASSWORD: [MessageHandler(filters.TEXT & ~filters.COMMAND, password)],
 			ACTION: [MessageHandler(filters.Regex("^(Enable input|Disable input|Lock screen|Format volumes|Play message)$"), action)],
 			FORMAT: [MessageHandler(filters.TEXT & ~filters.COMMAND, format)],
 			MESSAGE: [MessageHandler(filters.TEXT & ~filters.COMMAND, message)],