#!/usr/bin/env python3
import argparse
import getpass
import json
import os
import subprocess
import sys
import textwrap

SIGNAL_BUNDLEID = "org.whispersystems.signal"
SIGNAL_APPGROUP = "group.org.whispersystems.signal.group"
SIGNAL_APPGROUP_STAGING = "group.org.whispersystems.signal.group.staging"

SIGNAL_DEBUG_PAYLOAD_NAME = "dbPayload.txt"
SIGNAL_DEBUG_PAYLOAD_PASSPHRASE_KEY = "key"
SIGNAL_FALLBACK_DATABASE_PATH = "grdb/signal.sqlite"

DB_BROWSER_FOR_SQLITE_BUNDLEID = "net.sourceforge.sqlitebrowser"

quietMode=False
def failWithError(string):
    print("Error: " + string, file=sys.stderr)
    exit(1)

def printInfo(string = ""):
    if quietMode == False:
        print(string)

def runCommand(cmdList):
    result = subprocess.run(cmdList, text=True, capture_output=True)
    if result.returncode != 0:
        failWithError("Failed to run \"" + " ".join(cmdList) + "\". Status: " + str(result.returncode) + "\n" + result.stderr)
    return result.stdout


class Simulator:
    def __init__(self, searchString, useStaging):

        # Get JSON list of simulators matching searchString
        cmd = "xcrun simctl list -j devices " + searchString
        resultString = runCommand(cmd.split())
        simDict = json.loads(resultString)
        devicesByRuntime = simDict["devices"]

        # Parse all candidates
        candidates = []
        for runtime, devices in devicesByRuntime.items():
            os = self.parseOSFromRuntime(runtime)
            for device in devices:
                udid = device.get("udid")
                rawDevice = device.get("deviceTypeIdentifier")
                name = device.get("name")
                if udid != None:
                    deviceType = self.parseDeviceTypeFromRaw(rawDevice)
                    candidates.append({"os": os, "type": deviceType, "udid": udid, "name": name})

        # Select a candidate
        selectedCandidate = None

        if len(candidates) == 0:
            failWithError("Could not find a \"" + searchString + "\" simulator")
        elif len(candidates) == 1:
            selectedCandidate = candidates[0]
        else:
            if quietMode:
                failWithError("Multiple simulator candidates. Interactive selection not supported in quiet mode")
            for idx, candidate in enumerate(candidates):
                printInfo("{}:\t{:40}\t{} {} ({})".format(idx, candidate["name"], candidate["type"], candidate["os"], candidate["udid"]))

            while selectedCandidate == None:
                try:
                    idx = int(input("Select a simulator: "))
                    selectedCandidate = candidates[idx]
                except (ValueError, IndexError):
                    pass

        self.udid = selectedCandidate["udid"]
        self.groupID = SIGNAL_APPGROUP_STAGING if useStaging else SIGNAL_APPGROUP
        self.groupContainer = self.fetchGroupContainer(self.udid, self.groupID)
        printInfo("Selected simulator: " + selectedCandidate["name"] + " (" + selectedCandidate["udid"] + ")")
        printInfo("Using groupID: " + self.groupID)
        printInfo()

    def parseDebugPayload(self):
        path = self.groupContainer + "/" + SIGNAL_DEBUG_PAYLOAD_NAME
        try:
            fd = open(path, 'r')
            data = fd.read()
            payload = json.loads(data)
            return payload
        except IOError:
            return None

    def databasePath(self):
        return (self.groupContainer + "/" + SIGNAL_FALLBACK_DATABASE_PATH)

    def passphraseIfAvailable(self):
        debugPayload = self.parseDebugPayload()
        if debugPayload and SIGNAL_DEBUG_PAYLOAD_PASSPHRASE_KEY in debugPayload:
            return debugPayload[SIGNAL_DEBUG_PAYLOAD_PASSPHRASE_KEY]
        else:
            return None

    @staticmethod
    def parseOSFromRuntime(runtime):
        lastPeriodIdx = runtime.rfind('.')
        hypenatedOS = runtime[lastPeriodIdx+1:]
        return hypenatedOS.replace("-", ".")

    @staticmethod
    def parseDeviceTypeFromRaw(rawDevice):
        lastPeriodIdx = rawDevice.rfind('.')
        hypenatedOS = rawDevice[lastPeriodIdx+1:]
        return hypenatedOS.replace("-", " ")

    @staticmethod
    def fetchGroupContainer(udid, groupID):
        cmd = "xcrun simctl get_app_container {} {} {}".format(udid, SIGNAL_BUNDLEID, groupID)
        result = runCommand(cmd.split())
        return result.rstrip()

def preparePassphrase(passphrase):
    if len(passphrase) > 0 and passphrase[0] == 'x':
        return passphrase
    else:
        return "x'" + passphrase + "'"

def writeGuiEnvFile(passphrase, dbPath):
    dbName = os.path.basename(dbPath)
    envFilePath = os.path.join(os.path.dirname(dbPath), ".env")

    with open(envFilePath, "w", encoding="utf-8") as envFile:
        envFile.write(dbName + " = " + passphrase + "\n")
        envFile.write(dbName + "_plaintextHeaderSize = 32\n")

    return envFilePath

parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=textwrap.dedent('''\
                SQLCipher Command Line Interface

                    If providing a simulatorID (or accepting the default "Booted" simulator), passphrase retrieval
                    can be simplified by navigating to Signal Settings > Debug UI > Misc > Save plaintext database key.
                    If a database key could not be found and one was not provided through an argument, you'll be prompted
                    to enter one.

                    Alternatively, you can provide a sqlcipher path directly via command line arguments. In this case,
                    you'll be required to provide a database key through an argument or stdin.

                    If --use-gui is specified, this script will attempt to open the database using the "DB Browser for
                    SQLite" (DBBfS) application.

                    If --gui-auto-decrypt-with-plaintext-key is passed alongside --use-gui, the script will place the
                    passphrase in a file next to the database file such that DBBfS is able to automatically decrypt and
                    open the databse. Note that this file is in plaintext, and *ONLY USE* with databases containing
                    test data.
                '''),
        usage="%(prog)s [--simulator simID [--staging] | --path dbPath] [--passphrase passphrase] [--quiet] [--use-gui [--gui-auto-decrypt-with-plaintext-key]]")

group = parser.add_mutually_exclusive_group()
group.add_argument("--simulator", metavar="SIM", help="A string identifiying a simulator instance. (default: %(default)s).", default="booted")
group.add_argument("--path", help="Path to a sqlcipher DB")
parser.add_argument("--passphrase", metavar="PASS", help="The passphrase encrypting the database")
parser.add_argument("--staging", action='store_true', help="If a simulator is being targeted, specifies that the staging database should be used")
parser.add_argument("remainder", nargs=argparse.REMAINDER, metavar="--", help="All subsequent args will be interpreted as SQL. You probably want quotes here. Be careful with \"*\" since your shell will probably replace it. Ignored if using GUI")
parser.add_argument("--quiet", action='store_true', help="Suppress non-failing output")
parser.add_argument(
    "--use-gui",
    action='store_true',
    help="Tells the script to try and open DB Browser for SQLite"
)
parser.add_argument(
    "--gui-auto-decrypt-with-plaintext-key",
    action='store_true',
    help=(
        "Tells the script to try and have DB Browser for SQLite auto-decrypt the database by "
        "placing the key in plaintext next to the DB file. ONLY USE with DBs guaranteed to "
        "only contain test data"
    )
)
args = parser.parse_args()

quietMode=args.quiet
dbPath = None
passphrase = None

if args.path:
    dbPath = args.path
elif args.simulator:
    target = Simulator(args.simulator, args.staging)
    dbPath = target.databasePath()
    passphrase = target.passphraseIfAvailable()

if dbPath == None:
    failWithError("No valid database path")
elif os.path.isfile(dbPath) == False:
    failWithError("Not valid path " + dbPath)

if args.passphrase:
    passphrase = args.passphrase
if passphrase == None:
    passphrase = getpass.getpass("Please enter the passphrase. Alternatively, set up a plaintext database key in Debug UI > Misc > Save plaintext database key. Then, rerun the command. ")

if args.use_gui:
    if args.gui_auto_decrypt_with_plaintext_key:
        if passphrase == None or len(passphrase) == 0:
            failWithError("Missing sqlcipher passphrase for auto-decryption")

        passphrase = preparePassphrase(passphrase)
        envFilePath = writeGuiEnvFile(passphrase, dbPath)

        printInfo("Warning: saved passphrase to " + envFilePath + " for auto-decryption.")
    else:
        printInfo(textwrap.dedent('''\
            When prompted for the passphrase, select the SQLCipher 4 default settings.
            Then, select "Custom" and set the "Plaintext Header Size" to 32 from 0.
            Finally, select "Raw Key" instead of "Passphrase", manually enter "0x", and paste the key.
        '''))

    runCommand(["open", "-b", DB_BROWSER_FOR_SQLITE_BUNDLEID, dbPath])
else:
    if passphrase == None or len(passphrase) == 0:
        failWithError("No valid sqlcipher passphrase")

    passphrase = preparePassphrase(passphrase)

    sqlArgs = args.remainder
    if len(sqlArgs) > 0 and sqlArgs[0] == "--":
        sqlArgs.pop(0)
    sqlArgString = " ".join(sqlArgs)

    allArgs = [
        "sqlcipher",
        "-cmd", "PRAGMA key = \"" + passphrase + "\";",
        "-cmd", "PRAGMA cipher_plaintext_header_size = 32;",
        dbPath
    ]
    if len(sqlArgString) > 0:
        allArgs.append(sqlArgString)

    os.execvp("sqlcipher", allArgs)