From 4287d2b3775e1a1cc5f609ccf6b3eb5fe53f8c0f Mon Sep 17 00:00:00 2001 From: Andrew Kozlik Date: Sun, 3 Nov 2024 10:59:08 +0100 Subject: [PATCH] refactor(core): Add sanitize_input_script_type(). --- core/src/apps/bitcoin/common.py | 8 ++++++++ core/src/apps/bitcoin/get_ownership_id.py | 7 +------ core/src/apps/bitcoin/get_ownership_proof.py | 7 +------ 3 files changed, 10 insertions(+), 12 deletions(-) diff --git a/core/src/apps/bitcoin/common.py b/core/src/apps/bitcoin/common.py index 92f63be59b..3fb2054a16 100644 --- a/core/src/apps/bitcoin/common.py +++ b/core/src/apps/bitcoin/common.py @@ -190,6 +190,14 @@ def input_is_external_unverified(txi: TxInput) -> bool: ) +def sanitize_input_script_type(coin: CoinInfo, script_type: InputScriptType) -> None: + if script_type in SEGWIT_INPUT_SCRIPT_TYPES and not coin.segwit: + raise wire.DataError("Segwit not enabled on this coin.") + + if script_type == InputScriptType.SPENDTAPROOT and not coin.taproot: + raise wire.DataError("Taproot not enabled on this coin") + + def tagged_hashwriter(tag: bytes) -> HashWriter: from trezor.crypto.hashlib import sha256 from trezor.utils import HashWriter diff --git a/core/src/apps/bitcoin/get_ownership_id.py b/core/src/apps/bitcoin/get_ownership_id.py index dec6af5619..24247285ac 100644 --- a/core/src/apps/bitcoin/get_ownership_id.py +++ b/core/src/apps/bitcoin/get_ownership_id.py @@ -13,7 +13,6 @@ if TYPE_CHECKING: async def get_ownership_id( msg: GetOwnershipId, keychain: Keychain, coin: CoinInfo ) -> OwnershipId: - from trezor.enums import InputScriptType from trezor.messages import OwnershipId from trezor.wire import DataError @@ -34,11 +33,7 @@ async def get_ownership_id( if script_type not in common.INTERNAL_INPUT_SCRIPT_TYPES: raise DataError("Invalid script type") - if script_type in common.SEGWIT_INPUT_SCRIPT_TYPES and not coin.segwit: - raise DataError("Segwit not enabled on this coin") - - if script_type == InputScriptType.SPENDTAPROOT and not coin.taproot: - raise DataError("Taproot not enabled on this coin") + common.sanitize_input_script_type(coin, script_type) node = keychain.derive(msg.address_n) address = addresses.get_address(script_type, coin, node, msg.multisig) diff --git a/core/src/apps/bitcoin/get_ownership_proof.py b/core/src/apps/bitcoin/get_ownership_proof.py index 1bb2fe27c6..0e365b8616 100644 --- a/core/src/apps/bitcoin/get_ownership_proof.py +++ b/core/src/apps/bitcoin/get_ownership_proof.py @@ -19,7 +19,6 @@ async def get_ownership_proof( authorization: CoinJoinAuthorization | None = None, ) -> OwnershipProof: from trezor import TR - from trezor.enums import InputScriptType from trezor.messages import OwnershipProof from trezor.ui.layouts import confirm_action, confirm_blob from trezor.wire import DataError, ProcessError @@ -46,11 +45,7 @@ async def get_ownership_proof( if script_type not in common.INTERNAL_INPUT_SCRIPT_TYPES: raise DataError("Invalid script type") - if script_type in common.SEGWIT_INPUT_SCRIPT_TYPES and not coin.segwit: - raise DataError("Segwit not enabled on this coin") - - if script_type == InputScriptType.SPENDTAPROOT and not coin.taproot: - raise DataError("Taproot not enabled on this coin") + common.sanitize_input_script_type(coin, script_type) node = keychain.derive(msg.address_n) address = addresses.get_address(script_type, coin, node, msg.multisig)