Simple Bitcoin main and test network key generator.

BitcoinKey.py

import secrets
import hashlib
from enum import Enum

import ecdsa
import base58


class Key:
    class Prefix(Enum):
        MainPrivate = b'\x80'
        MainPublic = b'\x00'
        TestPrivate = b'\xEF'
        TestPublic = b'\x6F'

    __slots__ = ('private', 'public')

    def __init__(self):
        self.private = self.gen_private()
        self.public = self.gen_public(self.private)

    @staticmethod
    def gen_private() -> bytes:
        return secrets.token_bytes(32)

    @staticmethod
    def gen_public(private: bytes) -> bytes:
        pub_raw = ecdsa.SigningKey.from_string(private, curve=ecdsa.SECP256k1).verifying_key
        return b'\x04' + pub_raw.to_string()  # 04 for uncompressed

    @staticmethod
    def wif(key: bytes, prefix: Prefix) -> bytes:
        prefixed = prefix.value + key
        sha1 = hashlib.sha256(prefixed).digest()
        sha2 = hashlib.sha256(sha1).digest()
        checksum = sha2[:4]
        complete = prefixed + checksum
        return base58.b58encode(complete)

    @staticmethod
    def p2pkh(public: bytes, prefix: Prefix) -> bytes:
        hash_sha = hashlib.sha256(public).digest()
        hash_ripe = hashlib.new('ripemd160', hash_sha).digest()
        return Key.wif(hash_ripe, prefix)


def print_key(keys: Key) -> None:
    print('Private Key:', keys.private)
    print('Public Key:', keys.public)
    print('Private WIF:', Key.wif(keys.private, Key.Prefix.MainPrivate))
    print('Public P2PKH:', Key.p2pkh(keys.public, Key.Prefix.MainPublic))


def vanity(start: bytes) -> Key:
    while True:
        keys = Key()
        addr = Key.p2pkh(keys.public, Key.Prefix.MainPublic)
        if addr[1:len(start) + 1] == start:
            return keys


def main():
    key = vanity(b'mi')
    print_key(key)


if __name__ == '__main__':
    main()