dump-all.py

#!/usr/bin/env python3

# This file is a part of panopticon-2.
#
# Copyright (c) 2019 Ian Burgwin
# This file is licensed under The MIT License (MIT).
# You can find the full license text in LICENSE.md in the root of this project.

import argparse
import base64
import os
import re
from configparser import ConfigParser
from functools import lru_cache

import psycopg2


@lru_cache()
def clean_filename(string):
    return re.sub(r'[/\\:*?"<>|\x00-\x1f]', '', string)


parser = argparse.ArgumentParser(description='Dump logs from the panopticon-2 database.')
parser.add_argument('--config', help='Config file to use.', default='config.ini')

args = parser.parse_args()

config = ConfigParser()
config.read(args.config)

print('Connecting...')
con = psycopg2.connect(config['database']['dsn'])
with con:
    cur = con.cursor()
    cur2 = con.cursor()
    print('Executing query...')
    cur.execute('SELECT * FROM guild_messages_with_edits')
    print('Dumping contents...')
    for guild_id, guild_name, channel_id, channel_name, message_id, created_at, edited_at, user_tag, final_content in cur:
        path = os.path.join(f'{clean_filename(guild_name)}-{guild_id}', f'#{clean_filename(channel_name)}-{channel_id}')
        os.makedirs(path, exist_ok=True)
        filename = f'{created_at.strftime("%Y-%m-%d")}.log'
        with open(os.path.join(path, filename), 'a', encoding='utf-8') as o:
            o.write(f'[{"E:" if edited_at else ""}'
                    f'{base64.b64encode(message_id.to_bytes(8, "little")).decode("utf-8")}] '
                    f'{created_at.strftime("[%H:%M:%S.%f]")} '
                    f'<{user_tag}> ' + final_content.replace('\n', '\n(newline) ') + '\n')
            cur2.execute('SELECT url FROM guild_attachments WHERE message_id = %s', (channel_id,))
            for url in cur2:
                o.write('(attachment) ' + url + '\n')
    print('Done!')