import base58
import ecdsa


class KeyExtension:
    def get_private_key_from_wif(self, wif: str) -> bytes:
        if len(wif) == 0 or wif is None:
            raise ValueError(f"Empty WIF private key: {wif}")

        decoded = base58.b58decode_check(wif)
        if len(decoded) != 34 or decoded[0] != 0x80 or decoded[-1] != 0x01:
            raise ValueError("Not decode WIF by base58")

        private_key = decoded[1:-1]
        return private_key

    def get_public_key(self, private_key: bytes):
        if len(private_key) == 0 or private_key is None:
            raise ValueError(f"Empty private key: {private_key}")

        return ecdsa.SigningKey.from_string(private_key, curve=ecdsa.SECP256k1).get_verifying_key().to_string()