From fae99f511801fbc7402c85e22fc144ce0d3aa09f Mon Sep 17 00:00:00 2001
From: Tomer Vilensky <tomerv@fireblocks.com>
Date: Mon, 15 Jun 2020 15:42:27 +0300
Subject: [PATCH] Support setting AML user ID

---
 src/fireblocks-sdk.ts | 45 +++++++++++++++++++++++++++++++++++++++++++
 src/types.ts          |  7 +++++--
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/src/fireblocks-sdk.ts b/src/fireblocks-sdk.ts
index dded072e..224c268e 100644
--- a/src/fireblocks-sdk.ts
+++ b/src/fireblocks-sdk.ts
@@ -1,3 +1,4 @@
+
 import { ApiClient } from "./api-client";
 import { ApiTokenProvider } from "./api-token-provider";
 import { IAuthProvider } from "./iauth-provider";
@@ -439,4 +440,48 @@ export class FireblocksSDK {
     public async deleteExternalWalletAsset(walletId: string, assetId: string): Promise<WalletAssetResponse> {
         return await this.apiClient.issueDeleteRequest(`/v1/external_wallets/${walletId}/${assetId}`);
     }
+
+    /**
+     * Sets a user ID to pass as reference when performing AML screening
+     * @param vaultAccountId The vault account ID
+     * @param userId The user ID to set
+     */
+    public async setAmlUserIdForVaultAccount(vaultAccountId: string, userId: string): Promise<OperationSuccessResponse> {
+        return await this.apiClient.issuePostRequest(`/v1/vault/accounts/${vaultAccountId}/set_aml_user_id`, { userId });
+    }
+
+    /**
+     * Sets a user ID to pass as reference when performing AML screening
+     * @param walletId The ID of the internal wallet
+     * @param userId The user ID to set
+     */
+    public async setAmlUserIdForInternalWallet(walletId: string, userId: string): Promise<OperationSuccessResponse> {
+        return await this.apiClient.issuePostRequest(`/v1/internal_wallets/${walletId}/set_aml_user_id`, { userId });
+    }
+
+    /**
+     * Sets a user ID to pass as reference when performing AML screening
+     * @param walletId The ID of the external wallet
+     * @param userId The user ID to set
+     */
+    public async setAmlUserIdForExternalWallet(walletId: string, userId: string): Promise<OperationSuccessResponse> {
+        return await this.apiClient.issuePostRequest(`/v1/external_wallets/${walletId}/set_aml_user_id`, { userId });
+    }
+
+    /**
+     * Sets a user ID to pass as reference when performing AML screening
+     * @param vaultAccountId The vault account ID
+     * @param assetId The ID of the asset
+     * @param address The address
+     * @param tag The XRP tag, or EOS memo
+     * @param userId The user ID to set
+     */
+    public async setAmlUserIdForAddress(vaultAccountId: string, assetId: string, address: string, tag?: string, userId?: string): Promise<OperationSuccessResponse> {
+        let addressId = address;
+        if (tag && tag.length > 0) {
+            addressId = `${address}:${tag}`;
+        }
+
+        return await this.apiClient.issuePostRequest(`/v1/vault/accounts/${vaultAccountId}/${assetId}/addresses/${addressId}/set_aml_user_id`, { userId });
+    }
 }
diff --git a/src/types.ts b/src/types.ts
index 3e7df9d1..fffd6e90 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -3,6 +3,7 @@ export interface VaultAccountResponse {
     name: string;
     hiddenOnUI: boolean;
     assets: AssetResponse[];
+    amlUserId?: string;
 }
 
 export interface AssetResponse {
@@ -38,6 +39,7 @@ export interface WalletContainerResponse {
     id: string;
     name: string;
     assets: WalletAssetResponse[];
+    amlUserId?: string;
 }
 
 export interface CreateTransactionResponse {
@@ -67,6 +69,7 @@ export interface DepositAddressResponse {
     tag?: string;
     description?: string;
     type: string;
+    amlUserId?: string;
 }
 export interface GenerateAddressResponse {
     address: string;
@@ -79,13 +82,13 @@ export interface TransactionArguments {
     destination?: DestinationTransferPeerPath;
     amount: number | string;
     operation?: TransactionOperation;
-    waitForStatus?: boolean;
     fee?: number;
     gasPrice?: number;
     note: string;
     cpuStaking?: number;
     networkStaking?: number;
     autoStaking?: boolean;
+    amlUserId?: string;
 }
 
 export interface ExchangeResponse {
@@ -199,7 +202,7 @@ export enum TransactionStatus {
      */
     CONFIRMED = "CONFIRMED",
     COMPLETED = "COMPLETED",
-    PENDING_AML_CHECKUP = "PENDING_AML_CHECKUP",
+    PENDING_AML_SCREENING = "PENDING_AML_SCREENING",
     PARTIALLY_COMPLETED = "PARTIALLY_COMPLETED",
     CANCELLING = "CANCELLING",
     CANCELLED = "CANCELLED",