Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implemented pickle format as shared database for both bots (closes #18 ) #20

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
235 changes: 129 additions & 106 deletions discord_bot.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
'''
This file contains the discord bot to be called from within telegram_bot.py
'''
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
This file contains the discord bot to be called from within telegram_bot.py.
DiscordBot instantiates a second telegram bot for sending out the notifications.
"""
import os, discord
from dotenv import load_dotenv
from telegram import Bot
Expand All @@ -14,153 +17,172 @@ class DiscordBot:

def __init__(self):
"""
Constructor of the class. Initializes certain instance variables
and checks if everything's O.K. for the bot to work as expected.
Constructor of the class. Initializes some instance variables.
"""
# instantiate Telegram bot to send out messages to users
TELEGRAM_TOKEN = os.getenv('TELEGRAM_BOT_TOKEN')
# Instantiate Telegram bot to send out messages to users
TELEGRAM_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
self.telegram_bot = Bot(TELEGRAM_TOKEN)
#self.telegram_bot = TelegramBot()
# set of Discord usernames & roles that the Discord bot is listening to
self.listening_to = {'handles': set(), 'roles': set()}
# dictionary {discord username: telegram id}
self.discord_telegram_map = {'handles': {}, 'roles': {}}
# dictionary {telegram id: {data}}
# Set of Discord usernames & roles that the Discord bot is listening to
self.listening_to = {"handles": set(), "roles": set()}
# Reverse lookup {"handles": {discord username: {telegram id, telegram id}}
self.discord_telegram_map = {"handles": {}, "roles": {}}
# Dictionary {telegram id: {data}}
self.users = dict()
# path to database
self.data_path = './data'
self.debug_code = int(os.getenv('DEBUG'))
# Path to shared database (data entry via telegram_bot.py)
self.data_path = "./data"
self.debug_code = int(os.getenv("DEBUG"))
self.client = None

def refresh_data(self):
'''
"""
Populates/updates users, listening_to, discord_telegram_map.
'''
self.users = read_pickle(self.data_path)['user_data']

# update sets of notification triggers where available
for v in self.users.values():
[self.listening_to['roles'].add(x) for x in v['roles']]
try:
discord_handle = v['discord_username']
except KeyError:
continue
self.listening_to['handles'].add(discord_handle)

# create set of all TG ids requesting notifications for this handle
if discord_handle not in self.discord_telegram_map['handles']:
self.discord_telegram_map['handles'][discord_handle] = set()
self.discord_telegram_map['handles'][discord_handle].add(v['telegram_id'])

# create set of all TG ids requesting notifications for each role
for role in v['roles']:
if role not in self.discord_telegram_map['roles']:
self.discord_telegram_map['roles'][role] = set()
self.discord_telegram_map['roles'][role].add(v['telegram_id'])

print('Data updated!')

def send_to_TG(self, telegram_user_id, msg):
'''
"""
# Reload database from file
self.users = read_pickle(self.data_path)["user_data"]
print('\n\nself.users -> ', self.users)
# Update sets of notification triggers and reverse lookups
for k, v in self.users.items():
TG_id = k
# Add Discord handles to set of notification triggers
if "discord handle" in v:
handle = v["discord handle"]
self.listening_to["handles"].add(handle)
# Add Discord handle to reverse lookup
if handle not in self.discord_telegram_map["handles"]:
self.discord_telegram_map["handles"][handle] = set()
self.discord_telegram_map["handles"][handle].add(TG_id)

# Add Discord roles to set of notification triggers
if "discord roles" in v:
roles = v["discord roles"]
self.listening_to["roles"].update(roles)
# Add Discord roles to reverse lookup
for role in roles:
if role not in self.discord_telegram_map["roles"]:
self.discord_telegram_map["roles"][role] = set()
self.discord_telegram_map["roles"][role].add(TG_id)
# Debug area
print("\nself.discord_bot.refresh_data() called successfully!\n")
print("Data updated from Pickle file:")
print(f"users: {self.users}")
print(f"listening_to: {self.listening_to}")
print(f"discord_telegram_map: {self.discord_telegram_map}")

async def send_to_TG(self, telegram_user_id, msg):
"""
Sends a message a specific Telegram user id.
Uses Markdown V1 for inline link capability.
'''
self.telegram_bot.send_message(
"""
await self.telegram_bot.send_message(
chat_id=telegram_user_id,
text=msg,
disable_web_page_preview=True,
parse_mode='Markdown'
parse_mode="Markdown"
)

def send_to_all(self, msg):
'''
Sends a message to all Telegram bot users.
'''
telegram_ids = read_pickle(self.data_path).keys()
for telegram_id in telegram_ids:
self.send_to_TG(telegram_id, msg)
async def send_to_all(self, msg):
"""Sends a message to all Telegram bot users."""
TG_ids = list(self.users.keys())
for _id in TG_ids:
await self.send_to_TG(_id, msg)

def get_roles(self, discord_username, guild_id=1031616432049496225):
'''
"""
Takes a Discord username, returns all roles set for user in current guild.
'''
"""
#user = client.fetch_user(user_id)
guild = self.client.get_guild(guild_id)
user = guild.get_member_named(discord_username)
roles = [role.name for role in user.roles]
print(f'\n\n\nDEBUG DISCORD: Got these roles: {roles}')

print(f"\n\nDEBUG DISCORD: Got these roles: {roles}") # Debug only
return roles

def run_bot(self):
'''
Actual logic of the bot is stored here.
'''
# update data to listen to at startup
def get_listening_to(self, TG_id):
"""
Takes a TG username, returns whatever this user gets notifications for
currently.
"""
_map = self.discord_telegram_map
handles_active = {k for k, v in _map["handles"].items() if TG_id in v}
roles_active = {k for k, v in _map["roles"].items() if TG_id in v}

return {"handles": handles_active, "roles": roles_active}


async def run_bot(self):
"""Actual logic of the bot is stored here."""
# Update data to listen to at startup
self.refresh_data()

# fire up discord client
# Fire up discord client
intents = discord.Intents.default()
intents.members = True
intents.message_content = True
self.client = discord.Client(intents=intents)
client = self.client

# actions taken at startup
# Actions taken at startup
@client.event
async def on_ready():
all_handles = self.listening_to['handles']
all_roles = self.listening_to['roles']
handles = self.listening_to["handles"]
roles = self.listening_to["roles"]
mentions_update = f"""
Notifications active for mentions of {all_handles} and {all_roles}.
Notifications active for mentions of {handles} and {roles}.
"""
print(f'{client.user.name} has connected to Discord!')
msg = 'Discord bot is up & running!'
self.send_to_all(msg)
self.send_to_TG(self.debug_code, mentions_update)
print(f"\n\n{client.user.name} has connected to Discord!\n\n")
msg = "Discord bot is up & running!"
await self.send_to_all(msg)
await self.send_to_TG(self.debug_code, mentions_update)

# actions taken on every new Discord message
# Actions taken on every new Discord message
@client.event
async def on_message(message):

# handle mentions: forward to TG as in the Discord->Telegram lookup
for username in self.listening_to['handles']:
# Debug area
print(f"\n\nDiscord message -> {message}\n\n")
print(f"\n\nmessage.mentions type -> {type(message.mentions)}\n\n")
print(f"\n\nmessage.mentions -> {message.mentions}\n\n")
print(f"\n\nmessage.mentions == []-> {message.mentions == []}\n\n")
# TODO: Check for an empty message.mentions here to skip all the rest

# Handle mentions: Forward to TG as specified in Discord->Telegram lookup
for username in self.listening_to["handles"]:
user = message.guild.get_member_named(username)

# Debug area
print('\nusername in listening:', username)
print('user from guild:', user)
print('message.mentions:', message.mentions)

if user in message.mentions:
telegram_ids = self.discord_telegram_map['handles'][username]

for telegram_id in telegram_ids:

if self.users[telegram_id]['alerts_active']:
author = message.author
guild_name = message.guild.name
alias = user.display_name
url = message.jump_url
print(f'\n{author} mentioned {username}:')
contents = '@'+alias+message.content[message.content.find('>')+1:]
header = f"Mentioned by {author} in {guild_name}:\n\n"
link = '['+contents+']'+'('+url+')'
out_msg = header+link
self.send_to_TG(telegram_id, out_msg)

# role mentions: forward to TG as in the Discord->Telegram lookup
for role in self.listening_to['roles']:
print(f"\n\nUSER IN MENTIONS\n\n")
TG_ids = self.discord_telegram_map["handles"][username]

for _id in TG_ids:
author, guild_name = message.author, message.guild.name
alias, url = user.display_name, message.jump_url
print(f"\n{author} mentioned {username}!\n") # Debug only
contents = "@"+alias+message.content[message.content.find(">")+1:]
header = f"\nMentioned by {author} in {guild_name}:\n\n"
#link = "["+contents+"]"+"("+url+")"
out_msg = header+contents
await self.send_to_TG(_id, out_msg)

# Role mentions: Forward to TG as in the Discord->Telegram lookup
for role in self.listening_to["roles"]:
# probably some getter for role is needed for equality of objects
if role in message.mentions:
telegram_ids = self.discord_telegram_map['roles'][role]
author = message.author
guild_name = message.guild.name
alias = user.display_name
url = message.jump_url
contents = '@'+alias+message.content[message.content.find('>')+1:]
TG_ids = self.discord_telegram_map["roles"][role]
author, guild_name = message.author, message.guild.name
alias, url = user.display_name, message.jump_url
contents = "@"+alias+message.content[message.content.find(">")+1:]
header = f"Message to {role} in {guild_name}:\n\n"
link = '['+contents+']'+'('+url+')'
link = "["+contents+"]"+"("+url+")"
out_msg = header+link
print(f'\n{author} mentioned {role}:')
print(f"\n{author} mentioned {role}!\n") # Debug only

for telegram_id in telegram_ids:
if self.users[telegram_id]['alerts_active']:
self.send_to_TG(telegram_id, out_msg)
for _id in TG_ids:
await self.send_to_TG(_id, out_msg)

# DONE: Have bot also check for mentioned roles
# TODO: Have bot listen to specified subset of channels only
Expand All @@ -169,5 +191,6 @@ async def on_message(message):



DISCORD_TOKEN = os.getenv('DISCORD_TOKEN')
client.run(DISCORD_TOKEN)
DISCORD_TOKEN = os.getenv("DISCORD_TOKEN")
#client.run(DISCORD_TOKEN)
await client.start(DISCORD_TOKEN)
14 changes: 10 additions & 4 deletions main.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,20 +3,26 @@

from telegram_bot import TelegramBot
from discord_bot import DiscordBot
import asyncio
import tracemalloc
tracemalloc.start()

tg_bot = TelegramBot()
# initialize bots & insert Discord bot instance into TG bot
disc_bot = DiscordBot()
tg_bot = TelegramBot(disc_bot)

# initialize bots & insert Discord bot instance into TG bot
# Initialize Telegram bot
# Discord bot gets initialized from within async framework of TG bot
tg_bot.run_bot()
tg_bot.set_discord_instance(disc_bot)
#disc_bot.run_bot()




# DONE: add botpic, about info, description
# DONE: implement Telegram frontend (button menu for data entry)
# DONE: function for users to delete their data!
# DONE: Prompt for discord handle, guild/channel to listen to
# TODO: Switch off logging before bot is 'released'
# TODO: Find out best way to run both bots within 1 thread asynchronously
# TODO: Implement some way to assert that bot is not offline
# (i.e. call some bot function every 5 minutes & check for a reply)
Loading