channelsave.py
· 3.3 KiB · Python
Raw
#!/usr/bin/env python3
import discord
from peewee import *
from argparse import ArgumentParser
client = discord.Client()
db = SqliteDatabase("logs.db")
args = None
class BaseModel(Model):
class Meta:
database = db
class Author(BaseModel):
discriminator = CharField()
id = CharField(primary_key=True)
username = CharField()
class Server(BaseModel):
name = CharField()
id = CharField(primary_key=True)
class Channel(BaseModel):
name = CharField()
server = ForeignKeyField(Server, related_name="channels")
id = CharField(primary_key=True)
class Message(BaseModel):
author = ForeignKeyField(Author, related_name="messages")
server = ForeignKeyField(Server, related_name="messages")
channel = ForeignKeyField(Channel, related_name="messages")
content = CharField()
id = CharField(primary_key=True)
timestamp = DateField()
db.connect()
try:
db.create_tables([Author, Server, Channel, Message])
except OperationalError:
pass
@client.event
async def on_ready():
print("Logged in as {} ({})".format(client.user.name, client.user.id))
if client.user.bot:
print(discord.utils.oauth_url(client.user.id, permissions=discord.Permissions(74752)))
target = client.get_channel(args.channel)
if not target:
print("Channel not found")
await client.logout()
return
x = 0
server = None
channel = None
authors = {}
cache = []
try:
async for logline in client.logs_from(target, limit=args.limit):
if logline.server:
if not server:
try:
server = Server.get(id=logline.server.id)
except Server.DoesNotExist:
server = Server.create(id=logline.server.id, name=logline.server.name)
if not channel:
try:
channel = Channel.get(id=logline.channel.id)
except Channel.DoesNotExist:
channel = Channel.create(id=logline.channel.id, name=logline.channel.name, server=server)
if not logline.author.id in authors:
try:
authors[logline.author.id] = Author.get(id=logline.author.id)
except Author.DoesNotExist:
authors[logline.author.id] = Author.create(id=logline.author.id, username=logline.author.name, discriminator=logline.author.discriminator)
cache.append(dict(author=authors[logline.author.id], channel=channel, server=server, content=logline.content, id=logline.id, timestamp=logline.timestamp))
x += 1
if len(cache) >= 50:
Message.insert_many(cache).upsert().execute()
cache = []
print(x, end="\r")
finally:
Message.insert_many(cache).upsert().execute()
print("Saved {} messages from {}".format(x, target.name))
await client.logout()
try:
parser = ArgumentParser()
parser.add_argument("channel")
parser.add_argument("--limit", default=10**15, type=int)
args = parser.parse_args()
with open(".bottoken") as f:
creds = f.read()
if " " in creds or "." in creds:
client.run(*creds.split(None, 1), bot=False)
else:
client.run(creds.strip())
finally:
db.close()
| 1 | #!/usr/bin/env python3 |
| 2 | import discord |
| 3 | from peewee import * |
| 4 | from argparse import ArgumentParser |
| 5 | |
| 6 | client = discord.Client() |
| 7 | db = SqliteDatabase("logs.db") |
| 8 | args = None |
| 9 | |
| 10 | class BaseModel(Model): |
| 11 | class Meta: |
| 12 | database = db |
| 13 | |
| 14 | class Author(BaseModel): |
| 15 | discriminator = CharField() |
| 16 | id = CharField(primary_key=True) |
| 17 | username = CharField() |
| 18 | |
| 19 | class Server(BaseModel): |
| 20 | name = CharField() |
| 21 | id = CharField(primary_key=True) |
| 22 | |
| 23 | class Channel(BaseModel): |
| 24 | name = CharField() |
| 25 | server = ForeignKeyField(Server, related_name="channels") |
| 26 | id = CharField(primary_key=True) |
| 27 | |
| 28 | class Message(BaseModel): |
| 29 | author = ForeignKeyField(Author, related_name="messages") |
| 30 | server = ForeignKeyField(Server, related_name="messages") |
| 31 | channel = ForeignKeyField(Channel, related_name="messages") |
| 32 | content = CharField() |
| 33 | id = CharField(primary_key=True) |
| 34 | timestamp = DateField() |
| 35 | |
| 36 | db.connect() |
| 37 | try: |
| 38 | db.create_tables([Author, Server, Channel, Message]) |
| 39 | except OperationalError: |
| 40 | pass |
| 41 | |
| 42 | @client.event |
| 43 | async def on_ready(): |
| 44 | print("Logged in as {} ({})".format(client.user.name, client.user.id)) |
| 45 | if client.user.bot: |
| 46 | print(discord.utils.oauth_url(client.user.id, permissions=discord.Permissions(74752))) |
| 47 | target = client.get_channel(args.channel) |
| 48 | if not target: |
| 49 | print("Channel not found") |
| 50 | await client.logout() |
| 51 | return |
| 52 | x = 0 |
| 53 | server = None |
| 54 | channel = None |
| 55 | authors = {} |
| 56 | cache = [] |
| 57 | try: |
| 58 | async for logline in client.logs_from(target, limit=args.limit): |
| 59 | if logline.server: |
| 60 | if not server: |
| 61 | try: |
| 62 | server = Server.get(id=logline.server.id) |
| 63 | except Server.DoesNotExist: |
| 64 | server = Server.create(id=logline.server.id, name=logline.server.name) |
| 65 | if not channel: |
| 66 | try: |
| 67 | channel = Channel.get(id=logline.channel.id) |
| 68 | except Channel.DoesNotExist: |
| 69 | channel = Channel.create(id=logline.channel.id, name=logline.channel.name, server=server) |
| 70 | if not logline.author.id in authors: |
| 71 | try: |
| 72 | authors[logline.author.id] = Author.get(id=logline.author.id) |
| 73 | except Author.DoesNotExist: |
| 74 | authors[logline.author.id] = Author.create(id=logline.author.id, username=logline.author.name, discriminator=logline.author.discriminator) |
| 75 | cache.append(dict(author=authors[logline.author.id], channel=channel, server=server, content=logline.content, id=logline.id, timestamp=logline.timestamp)) |
| 76 | x += 1 |
| 77 | if len(cache) >= 50: |
| 78 | Message.insert_many(cache).upsert().execute() |
| 79 | cache = [] |
| 80 | print(x, end="\r") |
| 81 | finally: |
| 82 | Message.insert_many(cache).upsert().execute() |
| 83 | print("Saved {} messages from {}".format(x, target.name)) |
| 84 | await client.logout() |
| 85 | |
| 86 | try: |
| 87 | parser = ArgumentParser() |
| 88 | parser.add_argument("channel") |
| 89 | parser.add_argument("--limit", default=10**15, type=int) |
| 90 | args = parser.parse_args() |
| 91 | with open(".bottoken") as f: |
| 92 | creds = f.read() |
| 93 | if " " in creds or "." in creds: |
| 94 | client.run(*creds.split(None, 1), bot=False) |
| 95 | else: |
| 96 | client.run(creds.strip()) |
| 97 | finally: |
| 98 | db.close() |