-
Notifications
You must be signed in to change notification settings - Fork 0
/
bot.py
executable file
·129 lines (110 loc) · 3.77 KB
/
bot.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import json
import os
import sys
import telepot
import time
from pony import orm
from telepot.loop import MessageLoop
from dotenv import load_dotenv
# Classes
db = orm.Database()
class Chat(db.Entity):
_table_ = 'chats'
id = orm.PrimaryKey(int, auto=True)
telegram_id = orm.Required(str, unique=True)
main = orm.Required(int)
active = orm.Required(int)
name = orm.Optional(str)
topics = orm.Set(lambda: Topic, table='topics_for_chats', column='topic_id')
class Topic(db.Entity):
_table_ = 'topics'
id = orm.PrimaryKey(int, auto=True)
name = orm.Required(str)
tags = orm.Set(lambda: Tag, table='tags_for_topics', column='tag_id')
chats = orm.Set(lambda: Chat, table='topics_for_chats', column='chat_id')
class Tag(db.Entity):
_table_ = 'tags'
id = orm.PrimaryKey(int, auto=True)
name = orm.Required(str, unique=True)
topics = orm.Set(lambda: Topic, table='tags_for_topics', column='topic_id')
# Functions
def handle(msg):
txt = ''
try:
if 'text' in msg:
txt = txt + msg['text']
elif 'caption' in msg:
txt = txt + msg['caption']
if msg['chat']['type'] in ['group', 'supergroup'] and msg['new_chat_participant']:
print('Message: ' + str(msg))
if str(msg['new_chat_participant']['id']) == BOT_ID:
chatId = str(msg['chat']['id'])
with orm.db_session:
chat = Chat.get(telegram_id = chatId)
if chat == None:
chat = Chat(telegram_id = chatId, name = msg['chat']['title'], main = 0, active = 1, topics = availableTopics)
elif msg['chat']['type'] == 'channel' and isAllowed(msg) and txt != '':
print('Message: ' + str(msg))
with orm.db_session:
chats = Chat.select()[:]
for chat in chats:
if shouldForward(chat, txt):
try:
bot.forwardMessage(chat.telegram_id, SOURCE, msg['message_id'])
except:
print('Error forwarding message to ', chat.telegram_id)
except KeyError:
print('Whoops! KeyError')
def isAllowed(msg):
try:
if str(msg['chat']['id']) == SOURCE:
return True
except KeyError:
print('Whoops! KeyError')
return False
@orm.db_session
def shouldForward(chat, text):
if chat.active == 0:
return False
if chat.main > 0:
return True
interestedTags = []
topics = Chat[chat.id].topics
for topic in topics:
interestedTags.extend(list(topic.tags))
for tag in interestedTags:
if text.find('#' + tag.name) >= 0:
return True
return False
@orm.db_session
def loadTopics():
return Topic.select()
# MAIN
# Load env variables
load_dotenv()
BOT_ID = os.getenv('BOT_ID')
TOKEN = os.getenv('BOT_TOKEN')
PASSWORD = os.getenv('ADMIN_PASSWORD')
SOURCE = os.getenv('SOURCE')
DB_HOST = os.getenv('DB_HOST')
DB_USER = os.getenv('DB_USER')
DB_PASSWORD = os.getenv('DB_PASSWORD')
DB_DATABASE = os.getenv('DB_DATABASE')
# Connect to database
if DB_HOST == '' or DB_USER == '' or DB_PASSWORD == '' or DB_DATABASE == '':
sys.exit('No DB_HOST, DB_USER, DB_PASSWORD or DB_DATABASE in environment')
db.bind(provider='mysql', host=DB_HOST, user=DB_USER, passwd=DB_PASSWORD, db=DB_DATABASE)
db.generate_mapping()
# Load available topics
availableTopics = loadTopics()
# Start the bot
if TOKEN == '' or PASSWORD == '' or BOT_ID == '' or SOURCE == '':
sys.exit('No TOKEN, PASSWORD, SOURCE or BOT_ID in environment')
bot = telepot.Bot(TOKEN)
MessageLoop(bot, handle).run_as_thread()
print('Listening ...')
# Keep the program running.
while 1:
time.sleep(10)