twitools/bottools/methods.py

import ast, dbtools, html, io, logging, moviepy.editor, os, PIL.Image, random, re, setuptools, string, bottools.strings, sys, telegram.ext, telegram, time, twitools, bottools.streaming, bottools.thread,  urllib.request, tweepy

logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.WARN)
logger = logging.getLogger(__name__)

def getTwo(message):
 try:
  return twitools.twoBotHelper(message.chat_id)
 except ValueError:
  message.reply_text(bottools.strings.noauth)
 except tweepy.error.TweepError as e:
  raise

def twoExceptions(e, message):
 text = {
  32: bottools.strings.badToken,
  36: bottools.strings.selfSpam,
  64: bottools.strings.accountSuspended,
  88: bottools.strings.rateLimit,
  89: bottools.strings.badToken,
  99: bottools.strings.badToken,
  130: bottools.strings.overload,
  131: bottools.strings.twitterError,
  161: bottools.strings.followLimit,
  185: bottools.strings.tweetLimit,
  186: bottools.strings.longTweet,
  187: bottools.strings.dupTweet,
  205: bottools.strings.rateLimit,
  226: bottools.strings.automatedTweet,
  271: bottools.strings.selfMute,
  272: bottools.strings.notMuted,
  323: bottools.strings.multipleGIFs,
  326: bottools.strings.accountLocked,
  354: bottools.strings.longTweet
 }.get(e.api_code, bottools.strings.twoFail)

 message.reply_text(text)

def silence(bot, update):
 pass


def callback(bot, update):
 args = update.callback_query.data.split()

 try:
  feature = commands[args[0][1:]]
  try:
   status = feature(bot, update, args[1:])
  except:
   status = feature(bot, update)
 except Exception as e:
  update.callback_query.message.reply_text(bottools.strings.unknownCommand)
  return

 if status:
  tweetMessage(status, update.callback_query.message.chat_id, bot, callback = update.callback_query)

def captionHelper(bot, update):
 if update.message.caption.startswith("@%s /"):
  args = update.message.caption.split()[1:]
 else:
  args = update.message.caption.split()

 try:
  if args[0].startswith("/"):
   try:
    feature = commands[args[0][1:]]
    try:
     feature(bot, update, args[1:])
    except:
     feature(bot, update)
   except:
    update.message.reply_text(bottools.strings.unknownCommand)
  else:
   raise ValueError()
 except:
  explicitTweet(bot, update, args)

def mentionHelper(bot, update):
 args = update.message.text.split()
 
 try:
  feature = commands[args[1][1:]]
  try:
   feature(bot, update, args[2:])
  except:
   feature(bot, update)
 except Exception as e:
  update.message.reply_text(bottools.strings.unknownCommand)

# Actual methods:
# ---------------

def start(bot, update):
 update.message.reply_text(bottools.strings.start % {"name": setuptools.botname()})

def fish(bot, update):
 dbtools.dbHelper().addFish(update.message.chat_id)
 update.message.reply_text(bottools.strings.fishThanks)

def getTweetID(tlid, cid):
 try:
  db = dbtools.dbHelper()
  db.executeQuery("SELECT tid FROM timelines WHERE nr = %i AND cid = %i;" % (int(tlid), int(cid)))
  return db.getNext()[0]
 except Exception:
  raise ValueError("No such tweet in timeline")

def toggleTweet(bot, update):
 try:
  update.message.reply_text(bottools.strings.toggleTweet % (bottools.strings.toggleTweetOn if dbtools.dbHelper().toggleTweet(update.message.chat_id) else bottools.strings.toggleTweetOff))
 except tweepy.error.TweepError as e:
  bottools.methods.twoExceptions(e, update.message)

def toggleConfirmations(bot, update):
 try:
  update.message.reply_text(bottools.strings.toggleConfirmations % (bottools.strings.toggleTweetOn if dbtools.dbHelper().toggleConfirmations(update.message.chat_id) else bottools.strings.toggleTweetOff))
 except tweepy.error.TweepError as e:
  bottools.methods.twoExceptions(e, update.message)
  
def unknown(bot, update):
 update.message.reply_text(bottools.strings.unknownCommand)

def makeMenu(buttons, columns = 2, header = None, footer = None):
 menu = [buttons[i:i + columns] for i in range(0, len(buttons), columns)]
 if header:
  menu.insert(0, header)
 if footer:
  menu.append(footer)
 return menu

# Authentication process

def shareLocation(bot, update):
 buttons = [
  telegram.KeyboardButton(text = bottools.strings.shareLocationAgree, request_location = True),
 ]

 rmo = telegram.ReplyKeyboardMarkup(makeMenu(buttons, 1))
 update.message.reply_text(bottools.strings.shareLocation, reply_markup = rmo)

def unsetLocation(bot, update):
 dbtools.dbHelper().storeLocation(update.message.chat_id, "NULL", "NULL")
 update.message.reply_text(bottools.strings.unsetLocation)

def noLocation(bot, update):
 update.message.reply_text(bottools.strings.noLocation, reply_markup=telegram.ReplyKeyboardRemove())

def storeLocation(bot, update):
 cid = update.message.chat_id
 lat = update.message.location.latitude
 lon = update.message.location.longitude

 dbtools.dbHelper().storeLocation(cid, lat, lon)

 update.message.reply_text(bottools.strings.storeLocation, reply_markup=telegram.ReplyKeyboardRemove())

def auth(bot, update):
 message = update.message or update.callback_query.message
 db = dbtools.dbHelper()
 cid = message.chat_id

 auth = tweepy.OAuthHandler(setuptools.cke(), setuptools.cse())
 url = auth.get_authorization_url()
 dbtools.dbHelper().storeToken(cid, auth.request_token)

 rmo = telegram.InlineKeyboardMarkup([[telegram.InlineKeyboardButton("Verify", switch_inline_query_current_chat = "/verify")]])
 message.reply_text(bottools.strings.auth % url, reply_markup=rmo)

def verify(bot, update, args):
 db = dbtools.dbHelper()
 cid = update.message.chat_id

 try:
  auth = tweepy.OAuthHandler(setuptools.cke(), setuptools.cse())
  auth.request_token = ast.literal_eval(db.getToken(cid))

  try:
   auth.get_access_token(args[0])
   db.storeUser(cid, auth.access_token, auth.access_token_secret)
   update.message.reply_text(bottools.strings.verify)

  except Exception as e:
   db.deleteToken(update.message.chat_id)
   update.message.reply_text(bottools.strings.verifyfail)

  found = None

  for a in dbtools.dbHelper().accounts(cid):
   try:
    if twitools.twObject(ato=a[0], ase=a[1]).whoami().strip("@").lower() == args[0].strip("@").lower():
     found = a
     break
   except:
    pass

  if found:
   dbtools.dbHelper().deleteUserByAto(found[0])

 except:
  update.message.reply_text(bottools.strings.verifyimp)

def unauth(bot, update):
 global mentionstreams

 ato = dbtools.dbHelper().ato(update.message.chat_id)

 if ato in mentionstreams:
  mentionstreams.pop(ato).disconnect()
  
 dbtools.dbHelper().deleteUser(update.message.chat_id)
 update.message.reply_text(bottools.strings.unauth % setuptools.url())

def switch(bot, update, args):
 message = update.message or update.callback_query.message
 cid = message.chat_id
 found = None

 for a in dbtools.dbHelper().accounts(cid):
  try:
   if twitools.twObject(ato=a[0], ase=a[1]).whoami().strip("@").lower() == args[0].strip("@").lower():
    found = a
    break
  except:
   pass

 if found:
  dbtools.dbHelper().setActive(cid, a[0])
  message.reply_text(bottools.strings.switch % bottools.methods.getTwo(message).whoami().strip("@"))

  try:
   return getTwo(message).getTweet(getTweetID(args[1], cid))
  except:
   pass
  
 else:
  message.reply_text(bottools.strings.noauth)

def accounts(bot, update):
 accounts = []

 for a in dbtools.dbHelper().accounts(update.message.chat_id):
  try:
   accounts += [twitools.twObject(ato = a[0], ase = a[1]).whoami().strip("@")]
  except:
   pass

 buttons = []

 for account in accounts:
  buttons += [telegram.InlineKeyboardButton("@%s" % account, callback_data = "/switch %s" % account)]

 buttons += [telegram.InlineKeyboardButton("Login with Twitter", callback_data = "/login")]

 rmo = telegram.InlineKeyboardMarkup(bottools.methods.makeMenu(buttons, 1))

 try:
  name = getTwo(update.message).whoami()
 except:
  name = bottools.strings.accountNobody

 update.message.reply_text(bottools.strings.accounts % name, reply_markup = rmo)

# User methods

def follow(bot, update, args):
 try:
  two = bottools.methods.getTwo(update.message)
  for user in args:
   two.api.create_friendship(screen_name = user)
 except tweepy.error.TweepError as e:
   bottools.methods.twoExceptions(e, update.message)

def unfollow(bot, update, args):
 try:
  two = bottools.methods.getTwo(update.message)
  for user in args:
   two.api.destroy_friendship(screen_name = user)
 except tweepy.error.TweepError as e:
  bottools.methods.twoExceptions(e, update.message)

# Tweet methods

def explicitTweet(bot, update, args, reply = None):
 try:
  two = bottools.methods.getTwo(update.message)

  if update.message.photo or update.message.document or update.message.video or update.message.sticker:
   fid = update.message.document.file_id if update.message.document else update.message.sticker.file_id if update.message.sticker else update.message.video.file_id if update.message.video else update.message.photo[-1].file_id
   path = bot.getFile(fid).file_path
   media = urllib.request.urlopen(path)
   mobj = io.BytesIO(media.read())

   filename = path.split("/")[-1]

   if filename.split(".")[-1].lower() == "webp":
    out = io.BytesIO()
    PIL.Image.open(mobj).convert('RGB').save(out, format="JPEG")
    filename = "%s.jpg" % filename.split(".")[0]

   if update.message.document and filename.split(".")[-1].lower() == "mp4":
    temp = ''.join(random.choice(string.ascii_letters + string.digits) for i in range(32))

    with open("tmp/%s.%s" % (temp, filename.split(".")[-1]), "wb") as f:
     f.write(mobj.getvalue())

    moviepy.editor.VideoFileClip("tmp/%s.%s" % (temp, filename.split(".")[-1])).resize(0.3).write_gif("tmp/%s.gif" % temp)

    filename = "%s.gif" % temp
    out = open("tmp/%s.gif" % temp, "rb")

   else:
    out = mobj

   status = two.api.update_with_media(filename, ' '.join(args), in_reply_to_status_id=reply, file=out)
   out.close()

  else:
   try:
    status = two.tweet(' '.join(args), reply)
    bottools.methods.tweetMessage(status, update.message.chat_id, bot)
   except tweepy.error.TweepError as e:
    bottools.methods.twoExceptions(e, update.message)
   except Exception as e:
    logger.exception(e)

 except tweepy.error.TweepError as e:
  bottools.methods.twoExceptions(e, update.message)

def getMentions(tweet, two, args = []):
 sender = tweet.user.screen_name

 if "@%s" % sender.strip("@") in args:
  return []

 try:
  text = tweet.full_text
 except:
  text = tweet.text

 mentions = []

 for m in re.split('[^\w@]+', text):
  try:
   if m[0] == "@" and m[0].strip() != "" and not (m in mentions or m in args or m.strip() == "" or m == "@%s" % two.whoami().strip("@") or m == "@%s" % sender.strip("@")):
    mentions += [m]
  except:
   pass

 if not (sender.strip("@") == two.whoami().strip("@") or "@%s" % sender.strip("@") in mentions):
  mentions = ["@%s" % sender.strip("@")] + mentions

 return mentions

def reply(bot, update, args):
 try:
  reply = bottools.methods.getTweetID(args[0], update.message.chat_id)
  two = bottools.methods.getTwo(update.message)
  otweet = two.getTweet(reply)
  mentions = getMentions(otweet, two, args)
  pargs = mentions + args[1:]

 except:
  update.message.reply_text(bottools.strings.cantfind % args[0])
  raise

 bottools.methods.explicitTweet(bot, update, pargs, reply)

def quote(bot, update, args):
 try:
  reply = bottools.methods.getTweetID(args[0], update.message.chat_id)
  two = bottools.methods.getTwo(update.message)
  otweet = two.getTweet(reply)
  sender = otweet.user.screen_name

  args += ["https://twitter.com/%s/status/%i" % (sender, int(reply))]
 except:
  update.message.reply_text(bottools.strings.cantfind % args[0])
  raise
 
 bottools.methods.explicitTweet(bot, update, args[1:])

def retweet(bot, update, args):
 message = update.message or update.callback_query.message
 two = bottools.methods.getTwo(message)
 for tweet in args:
  try:
   tid = bottools.methods.getTweetID(tweet, message.chat_id)
   two.api.retweet(tid)
  except ValueError:
   message.reply_text(bottools.strings.cantfind % tweet)
  except tweepy.error.TweepError as e:
   bottools.methods.twoExceptions(e, message)

 if update.callback_query:
  time.sleep(0.5)
  return two.getTweet(tid)

def thread(bot, update, args):
 message = update.message or update.callback_query.message
 tid = bottools.methods.getTweetID(args[0], message.chat_id)
 two = getTwo(message)

 i = 0
 
 try:
  count = int(args[1])
 except:
  count = 5

 tweets = [two.getTweet(tid)]

 if tweets[-1].in_reply_to_status_id:
  nexttid = tweets[-1].in_reply_to_status_id

  while i < count:
   tweets += [two.getTweet(nexttid)]
   if not tweets[-1].in_reply_to_status_id:
    break
   nexttid = tweets[-1].in_reply_to_status_id
   i += 1

 tweets.reverse()

 for tweet in tweets:
  tweetMessage(tweet, message.chat_id, bot)

def like(bot, update, args):
 message = update.message or update.callback_query.message
 two = bottools.methods.getTwo(message)
 for tweet in args:
  try:
   tid = bottools.methods.getTweetID(tweet, message.chat_id)
   two.api.create_favorite(tid)
  except ValueError:
   message.reply_text(bottools.strings.cantfind % tweet)
  except tweepy.error.TweepError as e:
   bottools.methods.twoExceptions(e, message)

 if update.callback_query:
  time.sleep(0.5)
  return two.getTweet(tid)

def unlike(bot, update, args):
 message = update.message or update.callback_query.message
 two = bottools.methods.getTwo(message)
 for tweet in args:
  try:
   tid = bottools.methods.getTweetID(tweet, message.chat_id)
   two.api.destroy_favorite(tid)
  except ValueError:
   message.reply_text(bottools.strings.cantfind % tweet)
  except tweepy.error.TweepError as e:
   bottools.methods.twoExceptions(e, message)

 if update.callback_query:
  time.sleep(0.5)
  return two.getTweet(tid)

def tweet(bot, update):
 try:
  if dbtools.dbHelper().getTStatus(update.message.chat_id):
   bottools.methods.explicitTweet(bot, update, [update.message.text])
 except ValueError:
  update.message.reply_text(bottools.strings.noauth)
 except tweepy.error.TweepError as e:
  bottools.methods.twoExceptions(e, update.message)

# Timelines

def tweetMessage(status, cid, bot, force = False, callback = None, notified = None):
 db = dbtools.dbHelper()

 try:
  text = status.text
 except:
  text = status.full_text

 try:
  two = twitools.twoBotHelper(cid)
 except tweepy.error.TweepError as e:
  logger.exception(e)

 if not callback:
  if status.user.screen_name.strip("@") == two.whoami().strip("@") and not db.getCStatus(cid):
   return

 if not callback:
  try:
   db.executeQuery("SELECT MAX(nr) FROM timelines WHERE cid = %i;" % int(cid))
   i = int(db.getNext()[0]) + 1
  except:
   i = 1

  db.executeQuery("INSERT INTO timelines VALUES(%i, %i, %i);" % (cid, i, status.id))
  db.commit()

 else:
  i = int(callback.message.text.split()[1].strip(":"))

 buttons = []
 header = None

 if notified:
  if notified != db.ato(cid):
   ase = db.aseByAto(notified)
   oac = twitools.twObject(ato=notified, ase=ase).whoami()

   header = [telegram.InlineKeyboardButton(bottools.strings.messageSwitch % oac, callback_data = "/switch %s %i" % (oac, i))]

 if status.favorited:
  buttons += [telegram.InlineKeyboardButton("Unlike", callback_data = "/unlike %i" % i)]
 else:
  buttons += [telegram.InlineKeyboardButton("Like", callback_data = "/like %i" % i)]

 if status.retweeted:
  buttons += [telegram.InlineKeyboardButton("Retweeted", callback_data = "/silence")]
 else:
  buttons += [telegram.InlineKeyboardButton("Retweet", callback_data = "/retweet %i" % i)]

 buttons +=  [
  telegram.InlineKeyboardButton("Reply", switch_inline_query_current_chat = "/reply %i %s" % (i, ' '.join(getMentions(status, two)))),
  telegram.InlineKeyboardButton("Quote", switch_inline_query_current_chat = "/quote %i " % i)
 ]

 if status.in_reply_to_status_id:
  buttons += [telegram.InlineKeyboardButton("View Thread", callback_data = "/thread %i" % i)]

 try:
  rmu = telegram.InlineKeyboardMarkup(makeMenu(buttons, header=header))
 except Exception as e:
  logger.exception(e)

 if callback:
  bot.editMessageReplyMarkup(chat_id=callback.message.chat_id, message_id=callback.message.message_id, reply_markup=rmu)
 else:
  bot.sendMessage(chat_id = cid or callback.message.chat_id, text = "Tweet %i:\n%s (@%s) at %s:\n%s" % (i, status.author.name, status.author.screen_name, status.created_at, html.unescape(text)), reply_markup=rmu)

def trends(bot, update, args):
 try:
  count = int(args[0])
 except:
  count = 5

 two = getTwo(update.message)
 lat, lon = dbtools.dbHelper().getLocation(update.message.chat_id)

 lt = []

 try:
  if lat:
   woeid = two.api.trends_closest(lat, lon)[0]["woeid"]
  else:
   woeid = 1

  trends = two.api.trends_place(woeid)[0]['trends']

  buttons = []

  for trend in trends[:count]:
   buttons += [telegram.InlineKeyboardButton(trend['name'], callback_data = "/search %s" % trend['name'])]

  rmo = telegram.InlineKeyboardMarkup(makeMenu(buttons))

  update.message.reply_text(bottools.strings.trends, reply_markup = rmo)  

 except tweepy.error.TweepError as e:
  twoExceptions(e, update.message)

def search(bot, update, args):
 message = update.message or update.callback_query.message

 try:
  count = int(args[0])
  query = ' '.join(args[1:])
 except:
  count = 5
  query = ' '.join(args)

 try:
  two = bottools.methods.getTwo(message)

  lt = []

  for tweet in two.api.search(q=query, rpp=count, result_type="recent")[:count]:
   lt += [tweet]

  lt.reverse()

  for tweet in lt:
   tweetMessage(tweet, message.chat_id, bot)

 except tweepy.error.TweepError as e:
  bottools.methods.twoExceptions(e, message)

def user(bot, update, args):
 try:
  count = int(args[1])
 except:
  count = 5

 try:
  two = bottools.methods.getTwo(update.message)

  lt = []

  for status in two.api.user_timeline(args[0], count = count):
   lt += [status]

  lt.reverse()

  for status in lt:
   tweetMessage(status, update.message.chat_id, bot, True)

 except tweepy.error.TweepError as e:
  bottools.methods.twoExceptions(e, update.message)

def selfTweets(bot, update, args):
 try:
  count = int(args[0])
 except:
  count = 5

 bottools.methods.user(bot, update, [bottools.methods.getTwo(update.message).whoami(), count])

def timeline(bot, update, args = [10]):
 try:
  count = int(args[0])
 except:
  count = 10
   
 db = dbtools.dbHelper()
 db.executeQuery("DELETE FROM timelines WHERE cid = %i;" % int(update.message.chat_id))
 db.commit()
   
 try:
  two = bottools.methods.getTwo(update.message)
  lt = []
  
  for status in two.api.home_timeline(count=count):
   lt += [status]

  lt.reverse()

  for status in lt:
   tweetMessage(status, update.message.chat_id, bot, True)

 except tweepy.error.TweepError as e:
  bottools.methods.twoExceptions(e, update.message)

# Streaming

mentionstreams = {}

def makeStream(bot, cid, ato, ase):
# two = twitools.twObject(ato=ato, ase=ase)
# stream = tweepy.Stream(auth = two.auth, listener = bottools.streaming.BotStreamListener(bot, cid, ato))
# stream.filter(track=["%s" % two.whoami().strip("@")], async=True)
 stream = bottools.thread.BotStreamListener(bot, cid, ato)
 return stream
 
try:
 for u in dbtools.dbHelper().mentionsOn():
  mentionstreams[u[1]] = makeStream(telegram.Bot(token=setuptools.token()), u[0], u[1], u[2])
except Exception as e:
 print(e)

def mentionstream(bot, update):
 global mentionstreams

 try:
  message = update.message or update.callback_query.message
  db = dbtools.dbHelper()
  cid = message.chat_id

  db.toggleMentions(cid)
  ato = db.ato(cid)
  ase = db.ase(cid)

  if ato in mentionstreams:
   mentionstreams.pop(ato).disconnect()
   update.message.reply_text(bottools.strings.toggleMentions % bottools.strings.toggleTweetOff)
  else:
   mentionstreams[ato] = makeStream(bot, cid, ato, ase)
   update.message.reply_text(bottools.strings.toggleMentions % bottools.strings.toggleTweetOn)
 except tweepy.error.TweepError as e:
  bottools.methods.twoExceptions(e, update.message)

# Admin interaction

def togglebroadcasts(bot, update):
 try:
  update.message.reply_text(bottools.strings.toggleBroadcasts % (bottools.strings.toggleTweetOn if dbtools.dbHelper().toggleBroadcasts(update.message.chat_id) else bottools.strings.toggleTweetOff))
 except tweepy.error.TweepError as e:                                                                                                 
  bottools.methods.twoExceptions(e, update.message)

# Administrator

def isadmin(message):
 two = bottools.methods.getTwo(message)
 if two.whoami().strip("@") == setuptools.admin().strip("@"):
  return True
 return False

def restart(bot, update):
 global mentionstreams

 if bottools.methods.isadmin(update.message):
  update.message.reply_text(bottools.strings.restart)

  while len(mentionstreams) > 0:
   so = mentionstreams.pop(list(mentionstreams)[0])
   so.disconnect()
   
  os.execl(sys.executable, sys.executable, *sys.argv)

 else:
  bottools.methods.unknown(bot, update)

def broadcast(bot, update, args):
 if bottools.methods.isadmin(update.message):
  for u in dbtools.dbHelper().broadcastUsers():
   try:
    bot.sendMessage(chat_id = u, text = ' '.join(args))
   except:
    logger.exception("Could not send broadcast.")
 else:
  bottools.methods.unknown(bot, update)

def emergency(bot, update, args):
 if bottools.methods.isadmin(update.message):
  for u in dbtools.dbHelper().allUsers():
   try:
    bot.sendMessage(chat_id = u, text = ' '.join(args))
   except:
    logger.exception("Could not send emergency broadcast.")
 else:
  bottools.methods.unknown(bot, update)

# Variables

commands = {
 "accounts": accounts,
 "auth": auth,
 "broadcast": broadcast,
 "emergency": emergency,
 "fish": fish,
 "follow": follow,
 "help": start,
 "like": like,
 "login": auth,
 "logout": unauth,
 "mentionstream": mentionstream,
 "nolocation": noLocation,
 "quote": quote,
 "reply": reply,
 "restart": restart,
 "retweet": retweet,
 "search": search,
 "self": selfTweets,
 "sharelocation": shareLocation,
 "silence": silence,
 "start": start,
 "switch": switch,
 "thread": thread,
 "timeline": timeline,
 "togglebroadcasts": togglebroadcasts,
 "toggleconfirmations": toggleConfirmations,
 "togglementions": mentionstream,
 "toggletweet": toggleTweet,
 "trends": trends,
 "tweet": explicitTweet,
 "unauth": unauth,
 "unfollow": unfollow,
 "unlike": unlike,
 "unsetlocation": unsetLocation,
 "user": user,
 "verify": verify
}

pargs = [
 broadcast,
 emergency,
 follow,
 like,
 quote,
 reply,
 retweet,
 search,
 selfTweets,
 switch,
 thread,
 timeline,
 toggleTweet,
 trends,
 unlike,
 user,
 explicitTweet,
 unfollow,
 verify
]