From 7bb0092699299cc5fa403068dc91a590a62c5404 Mon Sep 17 00:00:00 2001
From: janniks <janniks@users.noreply.github.com>
Date: Wed, 23 Oct 2024 22:20:17 +0200
Subject: [PATCH] refactor: Change getStxAddress to use network parameter

---
 packages/transactions/src/keys.ts         |  8 ++---
 packages/wallet-sdk/src/models/account.ts | 36 ++++++++++++++++-------
 2 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/packages/transactions/src/keys.ts b/packages/transactions/src/keys.ts
index 4f8682dc1..e01bec80b 100644
--- a/packages/transactions/src/keys.ts
+++ b/packages/transactions/src/keys.ts
@@ -44,9 +44,9 @@ utils.hmacSha256Sync = (key: Uint8Array, ...msgs: Uint8Array[]) => {
 export function getAddressFromPrivateKey(
   /** Private key bytes or hex string */
   privateKey: PrivateKey,
-  network?: StacksNetworkName | StacksNetwork
+  network: StacksNetworkName | StacksNetwork = 'mainnet'
 ): string {
-  network = networkFrom(network ?? STACKS_MAINNET);
+  network = networkFrom(network);
   const publicKey = privateKeyToPublic(privateKey);
   return getAddressFromPublicKey(publicKey, network);
 }
@@ -55,9 +55,9 @@ export function getAddressFromPrivateKey(
 export function getAddressFromPublicKey(
   /** Public key bytes or hex string */
   publicKey: PublicKey,
-  network?: StacksNetworkName | StacksNetwork
+  network: StacksNetworkName | StacksNetwork = 'mainnet'
 ): string {
-  network = networkFrom(network ?? STACKS_MAINNET);
+  network = networkFrom(network);
   publicKey = typeof publicKey === 'string' ? hexToBytes(publicKey) : publicKey;
   const addrVer = addressHashModeToVersion(AddressHashMode.P2PKH, network);
   const addr = addressFromVersionHash(addrVer, hashP2PKH(publicKey));
diff --git a/packages/wallet-sdk/src/models/account.ts b/packages/wallet-sdk/src/models/account.ts
index bebea8bc1..e63703e57 100644
--- a/packages/wallet-sdk/src/models/account.ts
+++ b/packages/wallet-sdk/src/models/account.ts
@@ -4,13 +4,19 @@ import { HDKey } from '@scure/bip32';
 import { makeAuthResponse as _makeAuthResponse } from '@stacks/auth';
 import { bytesToHex, utf8ToBytes } from '@stacks/common';
 
+import { FetchFn, createFetchFn } from '@stacks/common';
 import {
   getPublicKeyFromPrivate,
   hashCode,
   hashSha256Sync,
   publicKeyToBtcAddress,
 } from '@stacks/encryption';
-import { createFetchFn, FetchFn } from '@stacks/common';
+import {
+  NetworkParam,
+  StacksNetwork,
+  StacksNetworkName,
+  TransactionVersion,
+} from '@stacks/network';
 import { getAddressFromPrivateKey } from '@stacks/transactions';
 import { connectToGaiaHubWithConfig, getHubInfo, makeGaiaAssociationToken } from '../utils';
 import { Account, HARDENED_OFFSET } from './common';
@@ -20,20 +26,28 @@ import {
   fetchProfileFromUrl,
   signAndUploadProfile,
 } from './profile';
-import { STACKS_MAINNET, STACKS_TESTNET, TransactionVersion } from '@stacks/network';
 
-export const getStxAddress = ({
+export function getStxAddress(
+  account: Account,
+  network?: StacksNetworkName | StacksNetwork
+): string;
+export function getStxAddress({
   account,
-  transactionVersion = TransactionVersion.Testnet,
+  network = 'mainnet',
 }: {
   account: Account;
-  transactionVersion?: TransactionVersion;
-}): string => {
-  return getAddressFromPrivateKey(
-    account.stxPrivateKey,
-    transactionVersion == TransactionVersion.Mainnet ? STACKS_MAINNET : STACKS_TESTNET // todo: refactor for `next` wallet update
-  );
-};
+} & NetworkParam): string;
+export function getStxAddress(
+  accountOrOptions: Account | ({ account: Account } & NetworkParam),
+  network: StacksNetworkName | StacksNetwork = 'mainnet'
+): string {
+  if ('account' in accountOrOptions) {
+    const { account, network = 'mainnet' } = accountOrOptions;
+    return getAddressFromPrivateKey(account.stxPrivateKey, network);
+  }
+
+  return getAddressFromPrivateKey(accountOrOptions.stxPrivateKey, network);
+}
 
 /**
  * Get the display name of an account.