#!/usr/bin/env python
#
# Electrum - lightweight Bitcoin client
# Copyright (C) 2015 Thomas Voegtlin
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
import os
import threading
import stat
import hmac
import base64
import zlib
from enum import IntEnum
from typing import Optional

from . import crypto
from .util import (profiler, InvalidPassword, WalletFileException, standardize_path,
                   test_read_write_permissions, os_chmod)

from .wallet_db import WalletDB
from .logging import Logger


def get_derivation_used_for_hw_device_encryption():
    return ("m"
            "/4541509'"      # ascii 'ELE'  as decimal ("BIP43 purpose")
            "/1112098098'")  # ascii 'BIE2' as decimal


class StorageEncryptionVersion(IntEnum):
    PLAINTEXT = 0
    USER_PASSWORD = 1
    XPUB_PASSWORD = 2


class StorageReadWriteError(Exception): pass


class StorageOnDiskUnexpectedlyChanged(Exception): pass


# TODO: Rename to Storage
class WalletStorage(Logger):

    # TODO maybe split this into separate create() and open() classmethods, to prevent some bugs.
    #      Until then, the onus is on the caller to check file_exists().
    def __init__(
        self,
        path,
        *,
        allow_partial_writes: bool = False,
    ):
        Logger.__init__(self)
        self.path = standardize_path(path)
        self._file_exists = bool(self.path and os.path.exists(self.path))
        self.logger.info(f"wallet path {self.path}")
        self._allow_partial_writes = allow_partial_writes
        # BTCQ BTQ1 format: pubkey now holds the 16-byte password verifier
        # (hex). The name is kept for compatibility with callers that copy
        # password state across storage instances (e.g. wallet.save_backup).
        self.pubkey = None  # verifier hex; truthy ⇔ wallet is unlocked
        self._password_salt: Optional[bytes] = None
        self._password_aes_key: Optional[bytes] = None
        self.decrypted = ''
        try:
            test_read_write_permissions(self.path)
        except IOError as e:
            raise StorageReadWriteError(e) from e
        if self.file_exists():
            with open(self.path, "rb") as f:
                self.raw = f.read().decode("utf-8")
                self.pos = f.seek(0, os.SEEK_END)
                self.init_pos = self.pos
            self._encryption_version = self._init_encryption_version()
        else:
            self.raw = ''
            self._encryption_version = StorageEncryptionVersion.PLAINTEXT
            self.pos = 0
            self.init_pos = 0

    def read(self):
        return self.decrypted if self.is_encrypted() else self.raw

    def write(self, data: str) -> None:
        try:
            mode = os.stat(self.path).st_mode
        except FileNotFoundError:
            mode = stat.S_IREAD | stat.S_IWRITE
        s = self.encrypt_before_writing(data)
        temp_path = "%s.tmp.%s" % (self.path, os.getpid())
        with open(temp_path, "wb") as f:
            try:
                os_chmod(temp_path, mode)  # set restrictive perms *before* we write data
            except PermissionError as e:  # tolerate NFS or similar weirdness?
                self.logger.warning(f"cannot chmod temp wallet file: {e!r}")
            f.write(s.encode("utf-8"))
            self.pos = f.seek(0, os.SEEK_END)
            f.flush()
            os.fsync(f.fileno())
        # assert that wallet file does not exist, to prevent wallet corruption (see issue #5082)
        if not self.file_exists():
            assert not os.path.exists(self.path)
        os.replace(temp_path, self.path)
        self._file_exists = True
        self.logger.info(f"saved {self.path}")

    def append(self, data: str) -> None:
        """ append data to file. for the moment, only non-encrypted file"""
        assert self._allow_partial_writes
        assert not self.is_encrypted()
        with open(self.path, "rb+") as f:
            pos = f.seek(0, os.SEEK_END)
            if pos != self.pos:
                raise StorageOnDiskUnexpectedlyChanged(f"expected size {self.pos}, found {pos}")
            f.write(data.encode("utf-8"))
            self.pos = f.seek(0, os.SEEK_END)
            f.flush()
            os.fsync(f.fileno())

    def _needs_consolidation(self):
        return self.pos > 2 * self.init_pos

    def should_do_full_write_next(self) -> bool:
        """If false, next action can be a partial-write ('append')."""
        return (
            not self.file_exists()
            or self.is_encrypted()
            or self._needs_consolidation()
            or not self._allow_partial_writes
        )

    def file_exists(self) -> bool:
        return self._file_exists

    def is_past_initial_decryption(self) -> bool:
        """Return if storage is in a usable state for normal operations.

        The value is True exactly
            if encryption is disabled completely (self.is_encrypted() == False),
            or if encryption is enabled but the contents have already been decrypted.
        """
        return not self.is_encrypted() or bool(self.pubkey)

    def is_encrypted(self) -> bool:
        """Return if storage encryption is currently enabled."""
        return self.get_encryption_version() != StorageEncryptionVersion.PLAINTEXT

    def is_encrypted_with_user_pw(self) -> bool:
        return self.get_encryption_version() == StorageEncryptionVersion.USER_PASSWORD

    def is_encrypted_with_hw_device(self) -> bool:
        return self.get_encryption_version() == StorageEncryptionVersion.XPUB_PASSWORD

    def get_encryption_version(self):
        """Return the version of encryption used for this storage.

        0: plaintext / no encryption
        1: BTQ1 — scrypt-derived AES-256-GCM, user password
           (BTCQ replaces upstream Electrum's BIE1/BIE2 ECIES path; the
           inherited XPUB_PASSWORD value is reserved but unused since
           hardware wallets are disabled on BTCQ.)
        """
        return self._encryption_version

    def _init_encryption_version(self):
        try:
            raw_bytes = base64.b64decode(self.raw, validate=True)
        except Exception:
            return StorageEncryptionVersion.PLAINTEXT
        magic = raw_bytes[:4]
        if magic == crypto.BTQ1_MAGIC:
            return StorageEncryptionVersion.USER_PASSWORD
        if magic in (b'BIE1', b'BIE2'):
            # Upstream Electrum (Bitcoin) wallet, not a BTCQ wallet. Fail
            # loudly rather than silently treat the file as plaintext —
            # otherwise the user gets a confusing "wallet appears empty"
            # state when the real problem is "you opened the wrong wallet".
            raise WalletFileException(
                "This wallet file uses upstream Electrum (Bitcoin) encryption "
                f"({magic.decode('ascii')}). BTCQ wallets use the BTQ1 format. "
                "If this is a Bitcoin Electrum wallet, open it in Electrum, "
                "not electrum-btcq."
            )
        return StorageEncryptionVersion.PLAINTEXT

    def _get_encryption_magic(self):
        v = self._encryption_version
        if v == StorageEncryptionVersion.USER_PASSWORD:
            return crypto.BTQ1_MAGIC
        if v == StorageEncryptionVersion.XPUB_PASSWORD:
            raise WalletFileException(
                "XPUB_PASSWORD storage encryption is not supported on BTCQ "
                "(hardware wallets are disabled). Use a user-password wallet."
            )
        raise WalletFileException('no encryption magic for version: %s' % v)

    def decrypt(self, password) -> None:
        """Raises an InvalidPassword exception on invalid password.

        On success, caches the per-wallet salt + AES key on the storage
        instance so subsequent encrypt_before_writing calls don't have to
        re-run scrypt.
        """
        if self.is_past_initial_decryption():
            return
        if self.raw:
            try:
                raw_bytes = base64.b64decode(self.raw, validate=True)
            except Exception as e:
                raise WalletFileException("wallet file is not valid base64") from e
            self._verify_btq1_magic(raw_bytes)
            salt = raw_bytes[4:4 + crypto.BTQ1_SALT_LEN]
            aes_key = crypto.btq1_derive_aes_key(password, salt)
            decrypted_bytes = crypto.btcq_storage_decrypt(
                password, raw_bytes, aes_key=aes_key,
            )
            s = zlib.decompress(decrypted_bytes).decode('utf8')
            self._password_salt = salt
            self._password_aes_key = aes_key
            self.pubkey = crypto.btq1_compute_verifier(salt, aes_key).hex()
        else:
            # Fresh empty file slated for encryption — caller will follow
            # up with set_password to initialise the salt/key.
            s = ''
        self.decrypted = s

    @staticmethod
    def _verify_btq1_magic(raw_bytes: bytes) -> None:
        magic = raw_bytes[:4]
        if magic != crypto.BTQ1_MAGIC:
            raise WalletFileException(
                f"unexpected encryption magic {magic!r}; expected BTQ1"
            )

    def encrypt_before_writing(self, plaintext: str) -> str:
        s = plaintext
        if self.pubkey:
            if self._password_salt is None or self._password_aes_key is None:
                raise WalletFileException(
                    "storage is unlocked but the BTQ1 salt/AES key are not "
                    "cached — decrypt() or set_password() must be called "
                    "before encrypt_before_writing()"
                )
            self.decrypted = plaintext
            data_bytes = bytes(plaintext, 'utf8')
            compressed = zlib.compress(data_bytes, level=zlib.Z_BEST_SPEED)
            blob = crypto.btcq_storage_encrypt(
                None,  # password unused — aes_key is supplied
                compressed,
                salt=self._password_salt,
                aes_key=self._password_aes_key,
            )
            s = base64.b64encode(blob).decode('ascii')
        return s

    def check_password(self, password: Optional[str]) -> None:
        """Raises an InvalidPassword exception on invalid password"""
        if not self.is_encrypted():
            if password is not None:
                raise InvalidPassword("password given but wallet has no password")
            return
        if not self.is_past_initial_decryption():
            self.decrypt(password)  # populates pubkey/salt/aes_key
            return
        assert self.pubkey is not None
        assert self._password_salt is not None
        # Fast path: re-derive verifier from password + cached salt and
        # compare against the cached verifier. Avoids re-decrypting the
        # whole file but still pays one scrypt for the password check.
        candidate_aes_key = crypto.btq1_derive_aes_key(password, self._password_salt)
        candidate_verifier = crypto.btq1_compute_verifier(self._password_salt, candidate_aes_key).hex()
        if not hmac.compare_digest(self.pubkey, candidate_verifier):
            raise InvalidPassword()

    def set_password(self, password, enc_version=None):
        """Set a password to be used for encrypting this storage.

        Generates a fresh salt and derives a new AES key — calling this
        with the same password twice produces different on-disk bytes,
        which is the desired behaviour for password-change flows.
        """
        if not self.is_past_initial_decryption():
            raise Exception("storage needs to be decrypted before changing password")
        if enc_version is None:
            enc_version = self._encryption_version
            if enc_version == StorageEncryptionVersion.PLAINTEXT and password:
                # First-time encryption on a previously plaintext wallet.
                enc_version = StorageEncryptionVersion.USER_PASSWORD
        if enc_version == StorageEncryptionVersion.XPUB_PASSWORD:
            raise WalletFileException(
                "XPUB_PASSWORD storage encryption is not supported on BTCQ "
                "(hardware wallets are disabled)."
            )
        if password and enc_version != StorageEncryptionVersion.PLAINTEXT:
            salt = os.urandom(crypto.BTQ1_SALT_LEN)
            aes_key = crypto.btq1_derive_aes_key(password, salt)
            self._password_salt = salt
            self._password_aes_key = aes_key
            self.pubkey = crypto.btq1_compute_verifier(salt, aes_key).hex()
            self._encryption_version = enc_version
        else:
            self.pubkey = None
            self._password_salt = None
            self._password_aes_key = None
            self._encryption_version = StorageEncryptionVersion.PLAINTEXT

    def copy_password_state_from(self, other: 'WalletStorage') -> None:
        """Copy the password-derived encryption state from another storage.

        Used by `wallet.save_backup` so the backup is encrypted with the
        same password as the source without re-prompting (or re-running
        scrypt). Both must already be past initial decryption.
        """
        self._encryption_version = other._encryption_version
        self.pubkey = other.pubkey
        self._password_salt = other._password_salt
        self._password_aes_key = other._password_aes_key

    def basename(self) -> str:
        return os.path.basename(self.path)

