#!/usr/bin/env python3

# This script is based on btcrecover extract-coinomi-privkey.py with changes
# similar to those Dhiru Kholia had made in multibit2john.py (yet different)

# extract-coinomi-privkey.py -- Coinomi private key extractor
# Copyright (C) 2014, 2015 Christopher Gurnee
#               2021 Stephen Rothery
#
# This file is part of btcrecover.
#
# btcrecover is free software: you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version
# 2 of the License, or (at your option) any later version.
#
# btcrecover is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see http://www.gnu.org/licenses/

import sys, os.path, binascii

# The protobuf parsing code is borrowed from btcrecover, and was
# initially automatically generated by Google's protoc utility.
# The precompiled file is shipped in protobuf/coinomi_pb2.py with John.
from protobuf.coinomi_pb2 import *

def process_file(filename):
    bname = os.path.basename(filename)
    try:
        f = open(filename, "rb")
        data = f.read()
    except IOError:
        e = sys.exc_info()[1]
        sys.stderr.write("%s\n" % str(e))
        return

    pb_wallet = Wallet()
    pb_wallet.ParseFromString(data)
    # print(pb_wallet)
    if pb_wallet.encryption_type == Wallet.UNENCRYPTED:
        raise ValueError("Coinomi wallet is not encrypted")
    if pb_wallet.encryption_type != Wallet.ENCRYPTED_SCRYPT_AES:
        raise NotImplementedError("Unsupported Coinomi wallet encryption type " + str(pb_wallet.encryption_type))
    if not pb_wallet.HasField("encryption_parameters"):
        raise ValueError("Coinomi wallet is missing its scrypt encryption parameters")

    # only need the final 2 encrypted blocks (half of it padding) plus the scrypt parameters
    encrypted_masterkey_part = pb_wallet.master_key.encrypted_data.encrypted_private_key[-32:]
    salt = pb_wallet.encryption_parameters.salt
    n = pb_wallet.encryption_parameters.n
    r = pb_wallet.encryption_parameters.r
    p = pb_wallet.encryption_parameters.p
    encrypted_data = binascii.hexlify(encrypted_masterkey_part).decode("ascii")
    salt = binascii.hexlify(salt).decode("ascii")
    sys.stdout.write("%s:$multibit$3*%s*%s*%s*%s*%s\n" % (bname, n, r, p, salt, encrypted_data))

    f.close()

if __name__ == "__main__":
    if len(sys.argv) < 2:
        sys.stderr.write("Usage: %s [Coinomi wallets files]\n" % sys.argv[0])
        sys.exit(1)

    for f in sys.argv[1:]:
        process_file(f)
