bot.py 2.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  1. #!/usr/bin/env python3.7
  2. import logging
  3. import queue
  4. import shlex
  5. import subprocess
  6. import sys
  7. from threading import Thread
  8. from ircclient import IRCClient, MODE_SEARCH, MODE_BOOK
  9. logging.basicConfig(level=logging.DEBUG)
  10. log = logging.getLogger(__name__)
  11. log.setLevel(logging.DEBUG)
  12. def usage():
  13. print("USAGE:")
  14. print("\t SEARCH <BOOK>")
  15. print("\t BOOK <BOT COMMAND>")
  16. def get_books_from_list(filename):
  17. f = open(filename, "r")
  18. ret = []
  19. for line in f.readlines():
  20. line = line.lower().strip()
  21. if not line.startswith('!') or not any(_type in line for _type in ['epub', 'mobi']):
  22. continue
  23. ret.append(line)
  24. log.info("Book matches: %s", line)
  25. ret = list(set(ret)) # dedup
  26. return ret
  27. def handle_files(files, mode):
  28. log.info("Unarchived files %s", files)
  29. out = []
  30. if mode == MODE_SEARCH:
  31. for f in files:
  32. if "searchbot" not in f.lower() and "searchook" not in f.lower():
  33. continue
  34. out.extend(get_books_from_list(f))
  35. return
  36. for f in files:
  37. if f.lower().endswith(".epub"):
  38. log.info("EPUB %s", f)
  39. new_fname = f.replace("epub", "mobi")
  40. p = subprocess.Popen(["ebook-convert", f, new_fname], stdout=subprocess.DEVNULL)
  41. # TODO log to file?
  42. p.wait()
  43. out.append(new_fname)
  44. out.append(f)
  45. # TODO make paths absolute?
  46. for filename in out:
  47. log.info(filename)
  48. def handle_results(q):
  49. while True:
  50. item = q.get()
  51. log.info("Got a result! %s", item)
  52. if item['type'] == 'files':
  53. handle_files(item['files'], item['mode'])
  54. def main():
  55. q = queue.Queue()
  56. rq = queue.Queue()
  57. client = IRCClient(command_queue=q, results_queue=rq)
  58. client.start()
  59. results_t = Thread(target=handle_results, args=(rq,))
  60. results_t.daemon = True
  61. results_t.start()
  62. for line in sys.stdin:
  63. line = line.strip().lower()
  64. try:
  65. split = shlex.split(line)
  66. except Exception as e:
  67. print(e)
  68. continue
  69. if split[0].lower() not in [MODE_SEARCH, MODE_BOOK]:
  70. usage()
  71. continue
  72. mode = split[0]
  73. query = " ".join(split[1:])
  74. data = {'query': query, 'mode': mode}
  75. q.put(data)
  76. client.join()
  77. results_t.join()
  78. if __name__ == "__main__":
  79. try:
  80. main()
  81. except KeyboardInterrupt:
  82. print("\nBye")