Bot
Woskoboyfrom datetime import timedelta, datetime as dtime
import re
from django.http import HttpResponseForbidden, HttpResponse
from django.views.decorators.csrf import csrf_exempt
from telebot import types
from tutugram.request_spider import CitiesKeyboard, TutuTrains # , YandexTrains
from tutugram.telegrambot import TBot, TOKEN
from tutugram.models import Chat
from tutugram.config import START_MESSAGE, SET_ARRIVAL, SEND_RUN, SET_DATES, ROUTE
from logging import getLogger
logger = getLogger(__name__)
DEPARTURE_STATE = True
ARRIVAL_STATE = False
DATE = 'date'
CODE = 'code'
DAYS = 10
TWICE_PASS = 1
bot = TBot()
@csrf_exempt
def index(request, token):
if token != TOKEN and request.method != "POST":
return HttpResponseForbidden('Invalid token')
start_update(request)
return HttpResponse('OK') # status=200)
def start_update(request):
try:
json_ = types.Update.de_json(request.body.decode('utf-8'))
if json_.message:
log(json_.message)
bot.process_new_messages([json_.message])
elif json_.callback_query:
log(json_.callback_query)
bot.process_new_callback_query([json_.callback_query])
elif json_.inline_query:
bot.process_new_inline_query([json_.inline_query])
else:
log(json_)
except Exception as e:
log('An error occurred while receiving JSON')
log(e)
@bot.message_handler(commands=['start', 'help'])
def echo_message(message):
bot.send_message(message.chat.id, START_MESSAGE)
reset_state(message.chat.id)
def reset_state(chat_id):
try:
db_chat = Chat.objects.get(id=chat_id)
db_chat.state = ARRIVAL_STATE
db_chat.save()
log('Reset to the initial state. Arrival state record to DB')
except Chat.DoesNotExist:
log('Chat DoesNotExist')
new_chat = Chat.objects.create(id=chat_id, date=dtime.now().strftime("%Y-%m-%d"))
new_chat.save()
@bot.message_handler(commands=['d'])
def echo_message(message):
log('Date are set manually')
chat_id = message.chat.id
try:
date = prepare_date(message.text[3:])
db_chat = Chat.objects.get(id=chat_id)
db_chat.date = date
db_chat.save()
bot.send_message(message.chat.id, SEND_RUN)
except IndexError:
log('Неправильно указана дата или введена команда')
except Chat.DoesNotExist:
log('DoesNotExist')
except Exception as e:
log('Проверьте введенные данные')
log(e)
def prepare_date(date):
return date if re.match('\d{4}-\d{2}-\d{2}', date) else dtime.now().strftime('%Y-%m-%d')
@bot.message_handler(commands=['run'])
def echo_message(message):
chat_id = message.chat.id
bot.send_chat_action(chat_id, 'typing')
create_route(chat_id)
reset_state(chat_id)
def create_route(chat_id):
try:
db_chat = Chat.objects.get(id=chat_id)
dep_code = db_chat.departure_code
arr_code = db_chat.arrival_code
date = db_chat.date
if dep_code and arr_code and date:
ya_url = ROUTE
ya_url += '&from={}&to={}&date={}'.format(dep_code, arr_code, date)
tutu_url = 'https://www.tutu.ru/poezda/rasp_d.php?nnst1={}&nnst2={}&date={}' \
.format(dep_code, arr_code, date)
log(tutu_url)
TutuTrains(bot, chat_id, tutu_url).run()
else:
bot.send_message(chat_id, 'Что-то пошло не так, попробуй ещё раз: /start')
log('Dep_code {}, Arr_code {}, Date {}'.format(dep_code, arr_code, date))
except Chat.DoesNotExist:
log('Bad route. Chat DoesNotExist')
@bot.message_handler(content_types=['text'])
def echo_message(message):
try:
db_chat = Chat.objects.get(id=message.chat.id)
change_state(db_chat)
except Chat.DoesNotExist:
log('Chat %s DoesNotExist ' % message.chat.id)
finally:
CitiesKeyboard(bot, message).run()
def change_state(db_chat):
if db_chat.state is DEPARTURE_STATE:
db_chat.state = ARRIVAL_STATE
log('State changed. Arrival state record to DB')
elif db_chat.state is ARRIVAL_STATE:
db_chat.state = DEPARTURE_STATE
log('State changed. Departure state record to DB')
db_chat.save()
@bot.callback_query_handler(func=lambda callback_query: True)
def callback_inline(callback_query):
chat_id = callback_query.message.chat.id
message_id = callback_query.message.message_id
data = callback_query.data
type_callback_data = get_type_data(data)
try:
db_chat = Chat.objects.get(id=chat_id)
if type_callback_data is CODE:
bot.edit_message_text('Код города: %s' % data, chat_id, message_id)
record_city_code(db_chat, data)
elif type_callback_data is DATE:
record_date(db_chat, data)
bot.edit_message_text('Выбрана дата %s' % data, chat_id, message_id)
bot.send_message(chat_id, SEND_RUN)
elif type_callback_data is None:
bot.send_message(chat_id, 'Что-то пошло не так, проверь дату, если задавал вручную'
'и попробуй ещё раз: /start')
log('Type callback is %s' % type_callback_data)
except Chat.DoesNotExist:
log('Chat %s Does Not Exist' % chat_id)
except Exception as e:
log(e)
def get_type_data(data):
if re.match('\d{7}', data):
return CODE
elif re.match('\d{4}-\d{2}-\d{2}', data):
return DATE
else:
return None
def record_city_code(db_chat, code):
count = db_chat.count
state = db_chat.state
if state is DEPARTURE_STATE:
db_chat.departure_code = code
log('Currently state is Departure. Departure code %s record to db' % code)
elif state is ARRIVAL_STATE:
db_chat.arrival_code = code
log('Currently state is Arrival. Arrival code %s record to db' % code)
if count is TWICE_PASS:
db_chat.count = 0
bot.send_message(db_chat.id, SET_DATES, reply_markup=make_keyboard(get_dates()))
else:
db_chat.count += 1
bot.send_message(db_chat.id, SET_ARRIVAL)
db_chat.save()
def record_date(db_chat, date):
db_chat.date = date
db_chat.save()
def get_dates():
return [(dtime.now() + timedelta(days=DAYS-i))
.strftime('%Y-%m-%d') for i in range(DAYS)]
def make_keyboard(dates):
keyboard = types.InlineKeyboardMarkup()
for date in dates:
keyboard.add(types.InlineKeyboardButton(text=date, callback_data=date))
return keyboard
def log(string):
logger.warning(string)