-
Notifications
You must be signed in to change notification settings - Fork 1
/
raspberrybot.py
196 lines (152 loc) · 6.39 KB
/
raspberrybot.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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
import json
import logging.config
import threading
from collections import defaultdict
import telebot
from raspberrybot import config, utils
from raspberrybot.models import TorrentStatus
INTERVAL = 60
EMOJI_THUMBS_UP = "\U0001F44D"
with open(config.LOGGING_CONFIG_FILE) as config_file:
logging.config.dictConfig(json.load(config_file))
log = logging.getLogger(__name__)
class TelegramBot(telebot.TeleBot):
def __init__(self, token, threaded=True):
super(TelegramBot, self).__init__(token, threaded)
log.info("Starting bot...")
self.downloading_torrents = set()
self._user_id = self.read_user_id()
if self._user_id:
log.info("Listening to user %d" % self.user_id)
self.send_init_message()
self.check_if_downloaded()
def send_init_message(self):
log.info("Sending init message")
self.send_message(self.user_id, "Raspberry Pi powered on")
@staticmethod
def read_user_id():
try:
with open(config.USER_ID_FILE, 'r') as user_id_file:
log.info("Reading user_id from file '%s'" % config.USER_ID_FILE)
return int(user_id_file.readline())
except IOError as e:
log.warn("Read user fail: %s" % e)
return None
@property
def user_id(self):
return self._user_id
@user_id.setter
def user_id(self, user_id):
log.info("Setting user_id to %s" % user_id)
self._user_id = user_id
with open(config.USER_ID_FILE, 'w') as user_id_file:
log.info("Writing user_id to file '%s'" % config.USER_ID_FILE)
user_id_file.write(str(self.user_id))
def validate_user(self, message):
return message.chat.id == self.user_id
def check_if_downloaded(self):
torrents = utils.torrent_list()
torrents_by_status = defaultdict(set)
for t in torrents:
if t.done == "100%":
torrents_by_status["done"].add(t)
else:
torrents_by_status["downloading"].add(t)
recently_completed = (self.downloading_torrents - torrents_by_status["downloading"]) & torrents_by_status["done"]
self.downloading_torrents = torrents_by_status["downloading"]
for t in recently_completed:
self.send_message(self._user_id, "Torrent '_%s_' has been downloaded " % t.name + EMOJI_THUMBS_UP, parse_mode="Markdown")
threading.Timer(INTERVAL, self.check_if_downloaded).start()
bot = TelegramBot(config.BOT_TOKEN)
@bot.message_handler(commands=['start'])
def send_welcome(message):
log.info("Handling command start")
if not bot.user_id:
log.info("New user")
bot.user_id = message.chat.id
bot.reply_to(message, "Added user %d" % bot.user_id)
elif message.chat.id == bot.user_id:
log.info("Previous user")
bot.reply_to(message, "Welcome back!")
else:
log.info("Invalid user %d" % message.chat.id)
bot.reply_to(message, "Invalid user. The Bot is linked with other user %s" % bot.user_id)
@bot.message_handler(commands=['public_ip'], func=bot.validate_user)
def get_public_ip(message):
log.info("Handling command get public ip - %s" % message.text)
bot.reply_to(message, utils.get_public_ip())
@bot.message_handler(commands=['local_ip'], func=bot.validate_user)
def get_local_ip(message):
log.info("Handling command get local ip - %s" % message.text)
bot.reply_to(message, utils.get_local_ip())
@bot.message_handler(commands=['uptime'], func=bot.validate_user)
def get_uptime(message):
log.info("Handling command uptime - %s" % message.text)
bot.reply_to(message, utils.get_uptime())
@bot.message_handler(commands=['torrent_add'], func=bot.validate_user)
def torrent_add(message):
log.info("Handling command add torrent - %s" % message.text)
args = message.text.split()
if len(args) < 2:
log.warn("Missing torrent URL")
bot.reply_to(message, "Missing torrent URL. Usage: /torrent_add [URL]")
return
url = args[1]
response = utils.torrent_add(url)
if "success" in response:
log.info("Torrent added")
bot.reply_to(message, "Torrent added")
else:
log.info("Error adding torrent")
bot.reply_to(message, response)
@bot.message_handler(commands=['torrent_list'], func=bot.validate_user)
def torrent_list_by_status(message):
log.info("Handling command list torrents - %s" % message.text)
torrents = utils.torrent_list()
log.info("Torrents: %s", torrents)
torrents_by_type = defaultdict(list)
for torrent in torrents:
torrents_by_type[torrent.status].append(torrent)
response = ""
for key in torrents_by_type.keys():
response += "*%s*\n" % TorrentStatus.conversion_to_string[key]
for t in torrents_by_type[key]:
response += t.format_telegram() + "\n"
bot.reply_to(message, response, parse_mode="Markdown")
@bot.message_handler(commands=['torrent_info'], func=bot.validate_user)
def torrent_info(message):
log.info("Handling command info torrent - %s" % message.text)
args = message.text.split()
if len(args) < 2:
log.warn("Missing torrent ID")
bot.reply_to(message, "Missing torrent ID. Usage: /torren_info [TORRENT_ID]")
return
torrent_id = args[1]
bot.reply_to(message, utils.torrent_info(torrent_id))
@bot.message_handler(commands=['torrent_remove'], func=bot.validate_user)
def torrent_remove(message):
log.info("Handling command remove torrent - %s" % message.text)
args = message.text.split()
if len(args) < 2:
log.warn("Missing torrent ID")
bot.reply_to(message, "Missing torrent ID. Usage: /torrent_remove [TORRENT_ID]")
return
torrent_id = args[1]
bot.reply_to(message, utils.torrent_remove(torrent_id))
@bot.message_handler(func=bot.validate_user)
def invalid_command(message):
log.warn("Invalid command - %s " % message.text)
bot.reply_to(message, "Invalid command '%s'" % message.text)
@bot.message_handler(func=lambda m: True)
def unauthorized_user(message):
log.warn("Unauthorized user - %s:%s " % (message.chat.id, message.text))
bot.reply_to(message, "User not allowed")
def polling():
# try:
bot.polling(none_stop=True)
# except Exception:
# log.exception("Error during polling. Waiting to reconnect...")
# time.sleep(5 * 60)
# polling()
if __name__ == "__main__":
polling()