betterbetterdiscord.py

import io
import os
import sys
import json
import struct
import shutil


def align_int(i, alignment):
    return i + (alignment - (i % alignment)) % alignment


class Asar:

    def __init__(self, filename, fp, header, base_offset):
        self.filename = filename
        self.fp = fp
        self.header = header
        self.base_offset = base_offset

    @classmethod
    def open(cls, filename):
        fp = open(filename, 'rb')

        # decode header
        data_size, header_size, header_object_size, header_string_size = struct.unpack('<4I', fp.read(16))

        base_offset = align_int(16 + header_string_size, 4)

        header_json = fp.read(header_string_size).decode('utf-8')
        header = json.loads(header_json)

        return cls(filename, fp, header, base_offset)

    @classmethod
    def from_folder(cls, path):
        offset = 0
        concatenated_files = b''

        def _folder_to_dict(path):
            nonlocal concatenated_files, offset
            result = {'files': {}}

            for f in os.scandir(path):
                if os.path.isdir(f.path):
                    result['files'][f.name] = _folder_to_dict(f.path)
                else:
                    size = f.stat().st_size
                    result['files'][f.name] = { 'size': size, 'offset': str(offset) }

                    with open(f.path, 'rb') as fp:
                        concatenated_files += fp.read()

                    offset += size

            return result

        header = _folder_to_dict(path)
        header_json = json.dumps(header, sort_keys=True, separators=(',', ':')).encode('utf-8')

        header_string_size = len(header_json)
        data_size = 4 # uint32
        aligned_size = align_int(header_string_size, data_size)
        header_size = aligned_size + 8
        header_object_size = aligned_size + data_size

        # pad with NULLs
        diff = aligned_size - header_string_size
        header_json = header_json + b'\0' * (diff) if diff else header_json

        fp = io.BytesIO()
        fp.write(struct.pack('<4I', data_size, header_size, header_object_size, header_string_size))
        fp.write(header_json)
        fp.write(concatenated_files)

        base_offset = align_int(16 + header_string_size, 4)

        return cls(path, fp, header, base_offset)

    def copy_extracted(self, source, destination):
        unpacked_dir = self.filename + '.unpacked'
        if not os.path.isdir(unpacked_dir):
            print("couldn't copy file {}, no extracted directory".format(source))
            return

        src = os.path.join(unpacked_dir, source)
        if not os.path.exists(src):
            print("couldn't copy file {}, doesn't exist".format(src))
            return

        dest = os.path.join(destination, source)
        shutil.copyfile(src, dest)

    def extract_file(self, source, info, destination):
        if 'offset' not in info:
            self.copy_extracted(source, destination)
            return

        offset = self.base_offset + int(info['offset'])
        self.fp.seek(offset)
        r = self.fp.read(int(info['size']))

        dest = os.path.join(destination, source)
        with open(dest, 'wb') as f:
            f.write(r)

    def extract_directory(self, source, files, destination):
        dest = os.path.normcase(os.path.join(destination, source))

        if not os.path.exists(dest):
            os.makedirs(dest)

        for name, info in files.items():
            item_path = os.path.join(source, name)

            if 'files' in info:
                self.extract_directory(item_path, info['files'], destination)
                continue

            self.extract_file(item_path, info, destination)

    def extract(self, destination):
        if os.path.exists(destination):
            raise FileExistsError()

        self.extract_directory('.', self.header['files'], destination)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.fp.close()

def main():
    # TODO: better interface for this lol
    # TODO: add argument for specifying where to find custom css file
    # TODO: use asar repacking code?
    #       (Asar.from_folder -> write Asar.from_folder.fp (BytesIO) to file)
    #       it's unnecessary (discord still works with app.asar as a folder)
    #       but maybe we'll need it at some point. shrug

    try:
        # TODO: find this in a better way
        discord_install_path = sys.argv[1]
    except IndexError:
        print('discord install path not given')
        return

    os.chdir(discord_install_path)

    custom_css_path = 'discord-custom.css'

    with Asar.open('./resources/app.asar') as a:
        a.extract('./resources/app')

    shutil.move('./resources/app.asar', './resources/original_app.asar')

    if not os.path.exists('./custom.css'):
        with open('./custom.css', 'w') as f:
            f.write('/* put your custom css here. */')

    # dumb
    with open('./resources/app/index.js', 'r') as f:
        entire_thing = f.read()

    # even dumber
    css_reload_script = """var customCSSPath = '{}';
    var customCSS = _fs2.default.readFileSync(customCSSPath, 'utf-8');
    mainWindow.webContents.on('dom-ready', function () {
      mainWindow.webContents.executeJavaScript(
        'window.myStyles = document.createElement("style");' +
        'window.myStyles.innerHTML = `' + customCSS + '`;' +
        'document.head.appendChild(window.myStyles);'
      );

      _fs2.default.watch(customCSSPath, { encoding: 'utf-8' }, function(eventType, filename) {
        if(eventType === 'change') {
          var changed = _fs2.default.readFileSync(customCSSPath, 'utf-8');
          mainWindow.webContents.executeJavaScript(
            "window.myStyles.innerHTML = `" + changed + "`;"
          );
        }
      });
    });""".format(custom_css_path)
    
    entire_thing = entire_thing.replace("mainWindow.webContents.on('dom-ready', function () {});", css_reload_script)

    with open('./resources/app/index.js', 'w') as f:
        f.write(entire_thing)


if __name__ == '__main__':
    main()