firefox/main.py

import os
import configparser
import shutil

from typing import Optional
from time import sleep
from argparse import ArgumentParser
from subprocess import run, Popen
from sys import stderr, stdout
from yaml import load, SafeLoader

from marionette_driver.addons import Addons
from marionette_driver.marionette import Marionette
from marionette_driver.geckoinstance import DesktopInstance
from mozprofile import FirefoxProfile


def shell_execute(cmd: str, env: Optional[dict] = None):
    if env is None:
        env = {}

    run(cmd, env=env, stdout=stdout, stderr=stderr, shell=True)


def profile_list() -> list:
    result = []

    ini = configparser.ConfigParser()
    ini.read(os.path.expanduser("~/.mozilla/firefox/profiles.ini"))
    for _, values in ini.items():
        for key, value in values.items():
            if key != "path":
                continue

            result.append(value)

    return list(set(result))


def profile_create(name_1: str):
    shell_execute(f"/usr/bin/firefox --createprofile {name_1} -headless")


def profile_configure(name_1: str, name_2: str):
    cwd = os.path.dirname(os.path.realpath(__file__))

    # Read config
    config = {}
    with open(f"{cwd}/profiles/{name_1}.yaml", "rt", encoding="utf-8") as handle:
        config = load(handle, Loader=SafeLoader)

    # Install signed addons first
    # Start firefox
    profile = FirefoxProfile(name_1)
    firefox = DesktopInstance(
        host="127.0.0.1",
        port=2828,
        bin="/usr/bin/firefox",
        app_args=["-remote-allow-system-access"],
        profile=profile,
        headless=1,
    )
    firefox.start()

    # Init marionette
    marionette = Marionette()
    marionette.start_session()

    # Install addons
    addons = Addons(marionette)
    for addon in config["addons"]:
        addons.install(f"{cwd}/addons/{addon['id']}-{addon['version']}.xpi")

    firefox.close()

    # Set prefs
    with open(f"{cwd}/{name_1}/user.js", "wt", encoding="utf-8") as handle:
        for key, value in config["prefs"].items():
            if type(value) == str:
                value = f'"{value}"'
            elif type(value) == bool:
                value = 'true' if value else 'false'
            else:
                value = value

            handle.write(f'user_pref("{key}", {value});\n')

    # Install profile
    from_path = f"{cwd}/{name_1}"
    to_path = os.path.expanduser(f"~/.mozilla/firefox/{name_2}")

    if os.path.exists(to_path):
        shutil.rmtree(to_path)

    shutil.copytree(from_path, to_path, dirs_exist_ok=True)
    shutil.rmtree(from_path)


def profile_init(name_1: str):
    proc = Popen(f"/usr/bin/firefox -p {name_1} -headless", shell=True)
    sleep(5)  # TODO: investigate why firefox take a time to apply extensions
    proc.kill()


def profile_patch(name_2: str):
    cwd = os.path.dirname(os.path.realpath(__file__))

    # Patch profile
    from_path = f"{cwd}/patches/search.json.mozlz4"
    to_path = os.path.expanduser(f"~/.mozilla/firefox/{name_2}/search.json.mozlz4")

    shutil.copy(from_path, to_path)


def main():
    parser = ArgumentParser()
    parser.add_argument("--configs", type=str, nargs="*", required=True)
    args = parser.parse_args()

    try:
        shutil.rmtree(os.path.expanduser("~/.mozilla"))
    except:
        pass

    names_1 = []

    for config in args.configs:
        names_1.append(config)

    for name_1 in names_1:
        profile_create(name_1)
        sleep(1)  # Prevent profile hash collision

    names_2 = profile_list()
    mapping = {}

    for name_1 in names_1:
        for name_2 in names_2:
            if name_1 in name_2:
                mapping[name_1] = name_2
                break

    for name_1, name_2 in mapping.items():
        profile_configure(name_1, name_2)
        profile_init(name_1)  # TODO: remove this workaround
        # profile_patch(name_2)  # TODO: patching not working for now, dynamic patching needed


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(e)