def disconnect(bot: Bot, update: Updater, user_data: dict):
"""
Closes the open SSH connection to the bridge computer.
After that, sends a 'Disconnected' message to the chat and returns to the
main menu.
Args:
bot (:obj:`telegram.bot.Bot`): The telegram bot instance.
update (:obj:`telegram.ext.update.Updater`): The Updater associated to
the bot.
user_data (:obj:`dict`): The dictionary with user variables.
"""
bridge_ip = Session.get_from(user_data).bridge_ip
# Save the new status of the computers
Session.get_from(user_data).computers.save()
# Close the connection
Session.get_from(user_data).end_connetion()
# Send the disconnect output
view.disconnect_output(bridge_ip).edit(update)
# Show the main menu again
menu.new_main(bot, update, user_data)
def job_update_computers_status(bot: Bot, job: Job):
user_data = job.context
if Session.get_from(user_data).connected:
update_computers_status(user_data)
else:
job.schedule_removal()
def __run_on_remote(self, computer: Computer, session: Session):
if not computer.on():
return (computer.name, computer.ip, "The computer is off", "")
session.copy_to_remote(computer.ip, self.bridge_path, self.remote_path)
# Replace the computer arguments (ip, mac...)
self.fill_computer_arguments(computer)
command = f"{self.remote_path} {' '.join(self.arguments.values())}"
return (
computer.name,
computer.ip,
(*session.run_on_remote(computer.ip, command, self.root)),
)
def get_username(bot: Bot, update: Updater, user_data: dict) -> int:
"""Get the username from the last messge. Ask for the password and wait"""
Session.get_from(user_data).username = update.message.text
view.ask_password().reply(update)
return PASSWORD
def get_password(bot: Bot, update: Updater, user_data: dict,
job_queue) -> ConversationHandler:
"""Get the password from the last message. End conversation"""
Session.get_from(user_data).password = update.message.text
general.connect(bot, update, user_data, job_queue)
return ConversationHandler.END
def confirm_connect_ip(bot: Bot, update: Updater, user_data: dict):
# Get the clicked ip
Session.get_from(user_data).bridge_ip = update.callback_query.data
# View
text = f"Connect to {Session.get_from(user_data).bridge_ip}"
view.yes_no(
text, yes_callback_data=State.GET_CREDENTIALS, no_callback_data=State.MAIN
).edit(update)
def select_department(bot: Bot, update: Updater, user_data: dict):
"""
Show the nodes on the first level of the :obj:`structure` field on the
:obj:`config.json` file.
Args:
bot (:obj:`telegram.bot.Bot`): The telegram bot instance.
update (:obj:`telegram.ext.update.Updater`): The Updater associated to
the bot.
user_data (:obj:`dict`): The dictionary with user variables.
"""
Session.get_from(user_data).route = []
route = Session.get_from(user_data).route
# View
view.structure(
route, states.config_file.get_sections(route), return_to=State.MAIN
).edit(update)
def execute_plugin(bot: Bot, update: Updater, user_data: dict):
session = Session.get_from(user_data)
plugin: Plugin = user_data["plugin"]
for name, ip, stdout, stderr in plugin.run(session):
view.plugin_output(name, ip, plugin.name, stdout,
stderr).reply(update, parse_mode=None)
menu.new_main(bot, update, user_data)
return ConversationHandler.END
def initialize_bridge(bot: Bot, update: Updater, user_data: dict):
session = Session.get_from(user_data)
plugin = Plugin("plugins/_bridge_initialization")
name, ip, stdout, stderr = next(plugin.run(session))
view.plugin_output(name,
ip,
"Bridge Initialization",
stdout,
stderr,
hide_header=True).reply(update)
def update_computers_status(user_data: dict):
session = Session.get_from(user_data)
plugin = Plugin("plugins/_computers_status")
_, _, stdout, _ = next(plugin.run(session))
# The stdout has the format `{ip} is {status}`
# TODO: Remove n^2 loop by adding dicts to the computers_json structure
for line in stdout.splitlines():
ip, _, status = line.split()
for computer in session.computers:
if ip == computer.ip:
computer.status = status
def run(self, session: Session):
session.copy_to_bridge(self.server_path, self.bridge_path, Plugin.COPY_MODE)
# Replace the session $arguments (username, password...)
self.fill_session_arguments(session)
# RUN ON BRIDGE
if self.source == PluginVar.SOURCE_BRIDGE:
command = f"{self.bridge_path} {' '.join(self.arguments.values())}"
yield "Bridge", session.bridge_ip, (
*session.run_on_bridge(command, root=self.root)
)
# RUN ON REMOTE
else:
with futures.ThreadPoolExecutor() as executor:
result_futures = [
executor.submit(self.__run_on_remote, computer, session)
for computer in session.computers.get_included_computers()
]
for future in futures.as_completed(result_futures):
yield future.result()
def ip_selection(bot: Bot, update: Updater, user_data: dict):
"""
Show a menu with the list of the computers that have a defined 'ip' field
in the corresponding .json file
"""
route = Session.get_from(user_data).route
# Get the clicked section
next_section = update.callback_query.data
route.append(next_section)
# Create a path to the .json file from the route
filepath = f"config/{'/'.join(route)}.json"
Session.get_from(user_data).computers = Computers(filepath)
Session.get_from(user_data).route = route
# View
view.ip_selection(
route,
Session.get_from(user_data).computers.get_computers(),
return_to=State.CONNECT,
).edit(update)
def main(bot: Bot, update: Updater, user_data: dict):
"""
Edit the last message to show the main menu. Depending on if the user is
connected or disconnected the view may change.
Args:
bot (:obj:`telegram.bot.Bot`): The telegram bot instance.
update (:obj:`telegram.ext.update.Updater`): The Updater associated to
the bot.
user_data (:obj:`dict`): The dictionary with user variables.
"""
session = Session.get_from(user_data)
# View
if not session.connected:
view.not_connected().edit(update)
else:
view.connected(session, Plugin.get_local_plugins()).edit(update)
def update_ips(bot: Bot, update: Updater, user_data: dict):
"""
Get all the macs and its associated ips from the local network. Then,
iterate through all the computers in :obj:`Computers`. If one of the
computer macs match with one of the local macs, update its associated ip to
the new value.
Args:
bot (:obj:`telegram.bot.Bot`): The telegram bot instance.
update (:obj:`telegram.ext.update.Updater`): The Updater associated to
the bot.
user_data (:obj:`dict`): The dictionary with user variables.
"""
session = Session.get_from(user_data)
# Get all the local ips for every local mac
plugin = Plugin("plugins/_local_arp_scan")
# Change view to executing
view.plugin_start(plugin.name).edit(update)
_, _, stdout, _ = next(plugin.run(session))
local_ips = {}
for line in stdout.splitlines():
ip, mac = line.strip().split()
local_ips[mac] = ip
for computer in session.computers.get_included_computers():
if computer.mac.lower() in local_ips:
last_ip = computer.ip
computer.ip = local_ips[computer.mac.lower()]
view.update_ip_output(computer, last_ip).reply(update)
session.computers.save()
menu.new_main(bot, update, user_data)
def connect(bot: Bot, update: Updater, user_data: dict, job_queue: JobQueue):
"""
Tries to open a SSH connection from the bot server to the bridge computer.
After that, sends a message to the chat with the result of the connection
(successed or failed) and returns to the main menu.
Args:
bot (:obj:`telegram.bot.Bot`): The telegram bot instance.
update (:obj:`telegram.ext.update.Updater`): The Updater associated to
the bot.
user_data (:obj:`dict`): The dictionary with user variables.
"""
session = Session.get_from(user_data)
# Try to connect to the client
session.start_connection()
# Send the status message
view.connect_output(session).reply(update)
if session.connected:
# Check if the bridge computer has all the required dependencies
initialize_bridge(bot, update, user_data)
# Check the status of each computers (which are alive or unreachable)
update_computers_status(user_data)
job_queue.run_repeating(
job_update_computers_status,
interval=states.config_file.check_status_interval,
context=user_data,
)
# Show the main menu again
menu.new_main(bot, update, user_data)
def start_plugin_from_download(bot: Bot, update: Updater, user_data: dict):
session = Session.get_from(user_data)
message = update.message
if not session.connected:
message.reply_text(
"You must be connected to a bridge computer before sending files!")
return ConversationHandler.END
# Download the file
file_object = message.document.get_file()
download_path = f"{states.config_file.server_tmp_dir}/{message.document.file_name}"
file_object.download(download_path)
# Make downloaded file executable by the user
os.chmod(download_path, 0o764)
user_data["plugin"] = Plugin(download_path)
view.plugin_start(user_data["plugin"].name).reply(update)
return collect_arguments(bot, update, user_data)
def exclude_computers(bot: Bot, update: Updater, user_data: dict):
"""
Change the attribute :obj:`included` for one/all computers to ``False``. If
a computer is not included, it won't be yielded in the loop while using the
:obj:`Computers.get_included_computers()` function.
Note:
The :obj:`query.data` value used by this function will be either the
string ``exclude-all`` for all computers, or ``exclude-[mac_address]``
for one computer.
Args:
bot (:obj:`telegram.bot.Bot`): The telegram bot instance.
update (:obj:`telegram.ext.update.Updater`): The Updater associated to
the bot.
user_data (:obj:`dict`): The dictionary with user variables.
"""
query = update.callback_query
# Extract the target from the string. Values: 'all' or '[mac_address]'
target = re.search(State.EXCLUDE_COMPUTERS, query.data).group(1)
computers = Session.get_from(user_data).computers
# Exclude all computers
if target == "all":
for computer in computers.get_computers():
computer.included = False
# Find the computer with the specified mac and exclude it
else:
computers.find(target).included = False
# Redraw the view
view.filter_computers(computers).edit(update)
def structure(bot: Bot, update: Updater, user_data: dict):
"""
Show the nodes on the level indicated by the route from the
:obj:`config.json` file.
Note:
The buttons shown by this function change dynamically depending on the
value of ``route``.
For example, if we have the following :obj:`config.json` file:
.. code-block:: python
{
"structure": [{
"name": "ESIT",
"sections": ["Industrial", "Ingeniería"]
},
{
"name": "MATFIS"
}]
}
Calling this function with :obj:`remote = []` will show the buttons
``ESIT`` and ``MATFIS``, but calling this function with :obj:`remote =
['ESIT']` will show the buttons ``Industrial`` and ``Ingeniería``, and
calling this function with :obj:`remote = ['MATFIS']` will show no
buttons.
Args:
bot (:obj:`telegram.bot.Bot`): The telegram bot instance.
update (:obj:`telegram.ext.update.Updater`): The Updater associated to
the bot.
user_data (:obj:`dict`): The dictionary with user variables.
"""
# Get the clicked section
next_section = update.callback_query.data
route = Session.get_from(user_data).route
# Remove section from the route if going backwards, otherwise append
if route and route[-1] == next_section:
route.pop()
else:
route.append(next_section)
# Return to the last section, or to the Connect menu if its the start oits
# the start of the route
return_to = ""
if len(route) <= 1:
return_to = State.CONNECT
else:
return_to = route[-1]
# Update the route
Session.get_from(user_data).route = route
# View
view.structure(
route, states.config_file.get_sections(route), return_to=return_to
).edit(update)
def filter_computers(bot: Bot, update: Updater, user_data: dict):
# View
computers = Session.get_from(user_data).computers
view.filter_computers(computers).edit(update)