A discord bot that uses markov chains to generate wacky messages
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

283 lines
8.7 KiB

  1. import discord
  2. from discord.ext import commands
  3. import requests
  4. from discord import Webhook, RequestsWebhookAdapter
  5. import os.path
  6. import asyncio
  7. import json
  8. import random
  9. import markovify
  10. import re
  11. from tqdm import tqdm
  12. intents = discord.Intents.default()
  13. intents.typing = False
  14. intents.presences = False
  15. intents.members = True
  16. # yoinked from nyandroid :3
  17. try:
  18. with open("token") as f:
  19. TOKEN = f.read().rstrip()
  20. except FileNotFoundError:
  21. print(
  22. "make sure you run this from the root dir, and that you have"
  23. " the token file present")
  24. exit(1)
  25. SOURCE_SERVER_ID = 689023661864386561
  26. TARGET_SERVER_ID = 767451816459108372
  27. TARGET_CHANNEL_ID = 768819426824028221
  28. CHANNEL_BLACKLIST = [
  29. # i have read the rules
  30. 689077340239560708,
  31. # mod-only channels
  32. 717341776339533894,
  33. 689052158674468905,
  34. # venting channels
  35. 724776502968975442,
  36. 689056634311409672,
  37. 697842470108659742,
  38. 689056648681357344,
  39. 716814102709665802,
  40. 692830672267771914,
  41. # other
  42. 690924986000736326,
  43. 715696652463112232,
  44. 695786643998900224,
  45. 689106682969718810,
  46. 689106643207716896
  47. ]
  48. MEMBER_WHITELIST = [
  49. 123960742754910211,
  50. 184444724520812545,
  51. 220918137569148928,
  52. 239787232666451980,
  53. 244284898184134657,
  54. 254310746450690048,
  55. 264264664219648000,
  56. 264970229514371072,
  57. 272864468386578432,
  58. 281626098943655937,
  59. 293040617200812053,
  60. 342097735866122242,
  61. 351350218953850883,
  62. 360543912982740992,
  63. 433790467830972417,
  64. 589057330138578966,
  65. 614213232017670147,
  66. 623131255331880963,
  67. 687740609703706630,
  68. 168487634866405376,
  69. 645736256663191553,
  70. 568909777401413657,
  71. 306257451638980619,
  72. 470350233419907129,
  73. 689053465598623845
  74. ]
  75. class Kommune2(commands.Bot):
  76. async def on_ready(self):
  77. print('Logged in as {0}!'.format(self.user))
  78. client = Kommune2(command_prefix='%', intents=intents)
  79. @client.event
  80. async def on_member_join(member):
  81. observer_role = client.get_guild(
  82. TARGET_SERVER_ID).get_role(767452466614108181)
  83. await member.add_roles(observer_role)
  84. @client.event
  85. async def on_message(message):
  86. if message.author == client.user:
  87. return
  88. else:
  89. await client.process_commands(message)
  90. # prevent the bot from counting commands as markov messages
  91. if message.content.startswith('%'):
  92. return
  93. if message.guild.id == TARGET_SERVER_ID:
  94. # pick a random user
  95. filename = random.choice(os.listdir("data/"))
  96. # load user data and messages
  97. f = open(f'data/{filename}')
  98. data = json.load(f)
  99. # make a newline-separated string of all messages
  100. messages = '\n'.join(data['messages'])
  101. # create webhook or find existing one
  102. webhooks = await message.channel.webhooks()
  103. if not len(webhooks):
  104. webhook = await message.channel.create_webhook(
  105. name='Text Generation Webhook')
  106. else:
  107. webhook = Webhook.from_url(
  108. url=webhooks[0].url,
  109. adapter=RequestsWebhookAdapter())
  110. if message.channel.id == TARGET_CHANNEL_ID:
  111. # ""think"" for a little bit before typing
  112. await asyncio.sleep(random.randint(5, 30))
  113. async with message.channel.typing():
  114. markov = markovify.NewlineText(messages, well_formed=False)
  115. out = markov.make_sentence(
  116. tries=200, max_overlap_ratio=0.7)
  117. # prevent empty message error
  118. if out is None:
  119. out = '[ could not generate a message, sowwy ]'
  120. # replace pings with names
  121. pattern = re.compile('<@!*(\d+)>')
  122. for m in re.finditer(pattern, out):
  123. userid = int(m.group(1))
  124. username = '**`@\u200b' + \
  125. client.get_user(userid).name + '`**'
  126. out = re.sub(
  127. pattern, username, out, 1)
  128. # oh h**eck i didn't think this would happen
  129. out = out.replace('@everyone', '**`@everyone`**')
  130. out = out.replace('@here', '**`@here`**')
  131. # simulate real typing speed
  132. #
  133. # average typing speed is around 200 chars per minute, which is ~3.3 per second -> it takes ~0.3s to type a single character
  134. await asyncio.sleep(len(out) * 0.3)
  135. # send the message using webhooks
  136. webhook.send(
  137. out,
  138. username=f"{data['displayname']}bot",
  139. avatar_url=data['avatar'])
  140. # delete error messages
  141. if message.content == '[ could not generate a message, sowwy ]':
  142. await message.delete()
  143. @client.command()
  144. async def populate(ctx, limit: int = 500):
  145. print(
  146. "Creating json files...\n"
  147. f"msg limit: {limit}"
  148. )
  149. # get the source guild to take messages from
  150. guild = discord.utils.find(
  151. lambda g: g.id == SOURCE_SERVER_ID, client.guilds)
  152. print(f'target server: {guild.name} [id: {guild.id}]\n')
  153. all_messages = []
  154. # save all messages and user ids to use them later
  155. with tqdm(guild.text_channels) as t:
  156. for channel in t:
  157. t.set_description(f'channel: {channel.name}')
  158. if channel.id not in CHANNEL_BLACKLIST:
  159. async for msg in channel.history(limit=limit):
  160. if msg.author.id in MEMBER_WHITELIST and len(msg.content):
  161. t.set_postfix(author=msg.author.display_name)
  162. message = {}
  163. message['id'] = msg.author.id
  164. message['text'] = msg.content
  165. all_messages.append(message)
  166. hivemind = {}
  167. hivemind['displayname'] = 'advnced bee brian'
  168. hivemind['avatar'] = 'https://upload.wikimedia.org/wikipedia/ru/6/63/Tool_-_Lateralus.jpg'
  169. hivemind['messages'] = []
  170. with tqdm(all_messages) as t:
  171. t.set_description('summoning the hivemind')
  172. for message in t:
  173. hivemind['messages'].append(message['text'])
  174. with open('data/hivemind.json', 'w') as out:
  175. json.dump(hivemind, out, indent=4)
  176. print('\n\nHIVEMIND GENERATED\n\n')
  177. # sort all messages into appropriate user files
  178. with tqdm(guild.members) as t:
  179. for member in t:
  180. t.set_description(f'member: {member.display_name}')
  181. # initialize the json structure
  182. data = {}
  183. data['displayname'] = member.display_name
  184. data['avatar'] = str(member.avatar_url)
  185. data['messages'] = []
  186. # set the user data
  187. if member.id in MEMBER_WHITELIST:
  188. with tqdm(all_messages) as t:
  189. for message in t:
  190. if message['id'] == member.id and len(message['text']):
  191. t.set_description(
  192. f"author: {guild.get_member(message['id']).display_name}")
  193. data['messages'].append(message['text'])
  194. # check if file exists
  195. filename = f'data/{member.id}.json'
  196. if os.path.isfile(filename):
  197. # replace the file if it exists
  198. os.remove(filename)
  199. with open(filename, 'w') as out:
  200. json.dump(data, out, indent=4)
  201. else:
  202. # otherwise, create a new file and dump data there
  203. with open(filename, 'w') as out:
  204. json.dump(data, out, indent=4)
  205. # clean up empty files
  206. with open(filename) as f:
  207. filedata = json.load(f)
  208. if not filedata.get('messages'):
  209. os.remove(filename)
  210. print("\n\nPOPULATION DONE\n\n")
  211. @client.command()
  212. async def update(ctx):
  213. # find the server to get data from
  214. guild = discord.utils.find(
  215. lambda g: g.id == SOURCE_SERVER_ID, client.guilds)
  216. # iterate over all user files
  217. for user in os.listdir("data/"):
  218. if re.match(r'\d+\.json', user):
  219. # load the data as a json object
  220. with open(f'data/{user}') as f:
  221. data = json.load(f)
  222. # remove `.json` from the filename
  223. user_id = int(user[:-5])
  224. # update their name and avatar
  225. data['displayname'] = guild.get_member(user_id).display_name
  226. data['avatar'] = str(guild.get_member(user_id).avatar_url)
  227. # write data
  228. with open(f'data/{user}', 'w') as f:
  229. json.dump(data, f)
  230. client.run(TOKEN)