Skip to content

Commit

Permalink
feat: new class for contracts
Browse files Browse the repository at this point in the history
  • Loading branch information
anshss committed Feb 3, 2025
1 parent adc2402 commit f00dd30
Show file tree
Hide file tree
Showing 13 changed files with 4,113 additions and 174 deletions.
1 change: 1 addition & 0 deletions packages/aw-signer/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
"@lit-protocol/constants": "7.0.2",
"@lit-protocol/contracts-sdk": "7.0.2",
"@lit-protocol/lit-node-client-nodejs": "7.0.2",
"@lit-protocol/misc": "^7.0.4",
"@lit-protocol/types": "7.0.2",
"bs58": "^6.0.0",
"ethers": "5.7.2",
Expand Down
1 change: 1 addition & 0 deletions packages/aw-signer/src/index.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
export { Admin } from './lib/admin';
export { Delegatee } from './lib/delegatee';
export { PkpToolRegistryContract } from './lib/contract';
export * from './lib/types';
export * from './lib/errors';
168 changes: 34 additions & 134 deletions packages/aw-signer/src/lib/admin.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@ import {
AdminConfig,
AgentConfig,
LitNetwork,
PkpInfo,
ToolInfoWithDelegateePolicy,
} from './types';
import {
Expand All @@ -19,30 +18,15 @@ import {
getRegisteredToolsAndDelegatees,
} from './utils/pkp-tool-registry';
import { LocalStorage } from './utils/storage';
import { loadPkpsFromStorage, mintPkp, savePkpsToStorage } from './utils/pkp';
import { AwSignerError, AwSignerErrorType } from './errors';
<<<<<<< HEAD
<<<<<<< HEAD

type AdminStorageLayout = {
[ethAddress: string]: {
privateKey: string;
pkps: PkpInfo[];
};
};

=======
import { getPkpEthAddressByTokenId } from './utils/pubkey-router';
>>>>>>> acb5096 (feat: pubkey router contract)
=======
import { getTokenIdByPkpEthAddress } from './utils/pubkey-router';
>>>>>>> 8ca8783 (checkpoint)
/**
* The `Admin` class is responsible for the ownership of the PKP (Programmable Key Pair),
* the registration and management of tools, policies, and delegatees.
*/
export class Admin {
private static readonly DEFAULT_STORAGE_PATH = './.law-signer-admin-storage';
private static readonly ADMIN_STORAGE_KEY = 'admins';
private static readonly DEFAULT_STORAGE_PATH = './.aw-signer-admin-storage';
// TODO: Add min balance check
// private static readonly MIN_BALANCE = ethers.utils.parseEther('0.001');

Expand Down Expand Up @@ -78,51 +62,11 @@ export class Admin {
this.adminWallet = adminWallet;
}

private static loadAdminsFromStorage(
storage: LocalStorage
): AdminStorageLayout {
const adminData = storage.getItem(Admin.ADMIN_STORAGE_KEY);
if (!adminData) {
return {};
}
return JSON.parse(adminData) as AdminStorageLayout;
}

private static saveAdminsToStorage(
storage: LocalStorage,
admins: AdminStorageLayout
): void {
storage.setItem(Admin.ADMIN_STORAGE_KEY, JSON.stringify(admins));
}

private static loadPkpsFromStorage(
storage: LocalStorage,
adminAddress: string
): PkpInfo[] {
const admins = Admin.loadAdminsFromStorage(storage);
return admins[adminAddress]?.pkps || [];
}

private static savePkpsToStorage(
storage: LocalStorage,
adminAddress: string,
pkps: PkpInfo[]
): void {
const admins = Admin.loadAdminsFromStorage(storage);
if (!admins[adminAddress]) {
admins[adminAddress] = { privateKey: '', pkps: [] };
}
admins[adminAddress].pkps = pkps;
Admin.saveAdminsToStorage(storage, admins);
}

private static async removePkpFromStorage(
storage: LocalStorage,
adminAddress: string,
pkpTokenId: string
) {
const admins = Admin.loadAdminsFromStorage(storage);
const pkps = admins[adminAddress]?.pkps || [];
const pkps = loadPkpsFromStorage(storage);
const index = pkps.findIndex((p) => p.info.tokenId === pkpTokenId);

if (index === -1) {
Expand All @@ -133,31 +77,7 @@ export class Admin {
}

pkps.splice(index, 1);
admins[adminAddress].pkps = pkps;
Admin.saveAdminsToStorage(storage, admins);
}

private static async mintPkp(
litContracts: LitContracts,
wallet: ethers.Wallet
): Promise<PkpInfo> {
const mintCost = await litContracts.pkpNftContract.read.mintCost();
if (mintCost.gt(await wallet.getBalance())) {
throw new AwSignerError(
AwSignerErrorType.INSUFFICIENT_BALANCE_PKP_MINT,
`${await wallet.getAddress()} has insufficient balance to mint PKP: ${ethers.utils.formatEther(
await wallet.getBalance()
)} < ${ethers.utils.formatEther(mintCost)}`
);
}

const mintMetadata = await litContracts.pkpNftContractUtils.write.mint();

return {
info: mintMetadata.pkp,
mintTx: mintMetadata.tx,
mintReceipt: mintMetadata.res,
};
savePkpsToStorage(storage, pkps);
}

/**
Expand Down Expand Up @@ -190,24 +110,22 @@ export class Admin {

let adminWallet: ethers.Wallet;
if (adminConfig.type === 'eoa') {
if (!adminConfig.privateKey) {
const storedPrivateKey = storage.getItem('privateKey');
const adminPrivateKey = adminConfig.privateKey || storedPrivateKey;

if (adminPrivateKey === null) {
throw new AwSignerError(
AwSignerErrorType.ADMIN_MISSING_PRIVATE_KEY,
'Admin private key not provided. Please provide a private key.'
'Admin private key not provided and not found in storage. Please provide a private key.'
);
}

adminWallet = new ethers.Wallet(adminConfig.privateKey, provider);

// Initialize storage for this admin if not already present
const admins = Admin.loadAdminsFromStorage(storage);
if (!admins[adminWallet.address]) {
admins[adminWallet.address] = {
privateKey: adminConfig.privateKey,
pkps: [],
};
Admin.saveAdminsToStorage(storage, admins);
// Only save if not already stored
if (!storedPrivateKey) {
storage.setItem('privateKey', adminPrivateKey);
}

adminWallet = new ethers.Wallet(adminPrivateKey, provider);
} else {
throw new AwSignerError(
AwSignerErrorType.ADMIN_MULTISIG_NOT_IMPLEMENTED,
Expand Down Expand Up @@ -242,15 +160,9 @@ export class Admin {
* Retrieves all PKPs stored in the Admin's (local) storage.
* @returns An array of PKP metadata.
*/
<<<<<<< HEAD
public async getPkps() {
return Admin.loadPkpsFromStorage(this.storage, this.adminWallet.address);
return loadPkpsFromStorage(this.storage);
}
=======
// public async getPkps() {
// return loadPkpsFromStorage(this.storage);
// }
>>>>>>> 8ca8783 (checkpoint)

/**
* Retrieves tokenId by pkpEthAddress
Expand All @@ -273,20 +185,10 @@ export class Admin {
* @throws If the PKP minting fails.
*/
public async mintPkp() {
<<<<<<< HEAD
const pkps = await this.getPkps();
const mintMetadata = await Admin.mintPkp(
this.litContracts,
this.adminWallet
);
pkps.push(mintMetadata);
Admin.savePkpsToStorage(this.storage, this.adminWallet.address, pkps);
=======
// const pkps = await this.getPkps();
const mintMetadata = await mintPkp(this.litContracts, this.adminWallet);
// pkps.push(mintMetadata);
// savePkpsToStorage(this.storage, pkps);s
>>>>>>> 8ca8783 (checkpoint)

return mintMetadata;
}
Expand Down Expand Up @@ -314,11 +216,7 @@ export class Admin {
);
}

await Admin.removePkpFromStorage(
this.storage,
this.adminWallet.address,
pkpTokenId
);
await Admin.removePkpFromStorage(this.storage, pkpTokenId);

return receipt;
}
Expand Down Expand Up @@ -420,7 +318,7 @@ export class Admin {
throw new Error('Tool policy manager not initialized');
}

const tx = await this.toolRegistryContract.enableTools(pkpTokenId, [
const tx = await this.toolRegistryContract.write.enableTools(pkpTokenId, [
toolIpfsCid,
]);

Expand All @@ -439,7 +337,7 @@ export class Admin {
throw new Error('Tool policy manager not initialized');
}

const tx = await this.toolRegistryContract.disableTools(pkpTokenId, [
const tx = await this.toolRegistryContract.write.disableTools(pkpTokenId, [
toolIpfsCid,
]);

Expand All @@ -461,7 +359,7 @@ export class Admin {
}

const [isRegistered, isEnabled] =
await this.toolRegistryContract.isToolRegistered(pkpTokenId, toolIpfsCid);
await this.toolRegistryContract.read.isToolRegistered(pkpTokenId, toolIpfsCid);

return { isRegistered, isEnabled };
}
Expand All @@ -478,7 +376,7 @@ export class Admin {
throw new Error('Tool policy manager not initialized');
}

const toolInfos = await this.toolRegistryContract.getRegisteredTools(
const toolInfos = await this.toolRegistryContract.read.getRegisteredTools(
pkpTokenId,
[toolIpfsCid]
);
Expand All @@ -497,7 +395,7 @@ export class Admin {
throw new Error('Tool policy manager not initialized');
}

const tools = await this.toolRegistryContract.getAllRegisteredTools(
const tools = await this.toolRegistryContract.read.getAllRegisteredTools(
pkpTokenId
);

Expand Down Expand Up @@ -536,7 +434,7 @@ export class Admin {
throw new Error('Tool policy manager not initialized');
}

return await this.toolRegistryContract.getDelegatees(pkpTokenId);
return await this.toolRegistryContract.read.getDelegatees(pkpTokenId);
}

/**
Expand All @@ -550,7 +448,7 @@ export class Admin {
throw new Error('Tool policy manager not initialized');
}

return await this.toolRegistryContract.isPkpDelegatee(
return await this.toolRegistryContract.read.isPkpDelegatee(
pkpTokenId,
ethers.utils.getAddress(delegatee)
);
Expand All @@ -567,9 +465,10 @@ export class Admin {
throw new Error('Tool policy manager not initialized');
}

const tx = await this.toolRegistryContract.addDelegatees(pkpTokenId, [
delegatee,
]);
const tx = await this.toolRegistryContract.write.addDelegatees(
pkpTokenId,
[delegatee]
);

return await tx.wait();
}
Expand All @@ -586,9 +485,10 @@ export class Admin {
throw new Error('Tool policy manager not initialized');
}

const tx = await this.toolRegistryContract.removeDelegatees(pkpTokenId, [
delegatee,
]);
const tx = await this.toolRegistryContract.write.removeDelegatees(
pkpTokenId,
[delegatee]
);

return await tx.wait();
}
Expand All @@ -610,7 +510,7 @@ export class Admin {
throw new Error('Tool policy manager not initialized');
}

const result = await this.toolRegistryContract.isToolPermittedForDelegatee(
const result = await this.toolRegistryContract.read.isToolPermittedForDelegatee(
pkpTokenId,
toolIpfsCid,
ethers.utils.getAddress(delegatee)
Expand Down Expand Up @@ -860,15 +760,15 @@ export class Admin {
* @param ipfsCid - The IPFS CID of the tool.
* @param delegatee - The address of the delegatee.
* @param parameterNames - An array of policy parameter names.
* @returns A promise that resolves to an array of tuples containing policy parameter names and values.
* @returns A promise that resolves to an array of policy parameter values.
* @throws If the tool policy registry contract is not initialized.
*/
public async getToolPolicyParametersForDelegatee(
pkpTokenId: string,
ipfsCid: string,
delegatee: string,
parameterNames: string[]
): Promise<{ name: string; value: string }[]> {
) {
if (!this.toolRegistryContract) {
throw new Error('Tool policy manager not initialized');
}
Expand Down
Loading

0 comments on commit f00dd30

Please sign in to comment.