From 286a491249eb95d93d3d324a8cd68c6fac67b81c Mon Sep 17 00:00:00 2001
From: Javier Viola <363911+pepoviola@users.noreply.github.com>
Date: Tue, 11 Jun 2024 11:17:14 +0200
Subject: [PATCH] =?UTF-8?q?feat:=20add=20config=20to=20force=20the=20decor?=
 =?UTF-8?q?ator=20to=20use=20and=20create=20a=20new=20one=20(=E2=80=A6=20(?=
 =?UTF-8?q?#1816)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* feat: add config to force the decorator to use and create a new one (generic-evm)

* typo
---
 .../src/chain-decorators/generic-evm.ts       | 91 +++++++++++++++++++
 .../src/chain-decorators/index.ts             | 15 ++-
 .../orchestrator/src/configGenerator.ts       |  1 +
 .../packages/orchestrator/src/sharedTypes.ts  |  1 +
 4 files changed, 107 insertions(+), 1 deletion(-)
 create mode 100644 javascript/packages/orchestrator/src/chain-decorators/generic-evm.ts

diff --git a/javascript/packages/orchestrator/src/chain-decorators/generic-evm.ts b/javascript/packages/orchestrator/src/chain-decorators/generic-evm.ts
new file mode 100644
index 000000000..d4bd03736
--- /dev/null
+++ b/javascript/packages/orchestrator/src/chain-decorators/generic-evm.ts
@@ -0,0 +1,91 @@
+import { Keyring } from "@polkadot/api";
+import { u8aToHex } from "@polkadot/util";
+import { cryptoWaitReady } from "@polkadot/util-crypto";
+import { CreateLogTable, decorators } from "@zombienet/utils";
+import {
+  GenesisNodeKey,
+  getRuntimeConfig,
+  readAndParseChainSpec,
+  writeChainSpec,
+} from "../chainSpec";
+import { generateKeyForNode as _generateKeyForNode } from "../keys";
+import { Node } from "../sharedTypes";
+
+async function generateKeyForNode(nodeName?: string): Promise<any> {
+  const keys = await _generateKeyForNode(nodeName);
+
+  await cryptoWaitReady();
+
+  const eth_keyring = new Keyring({ type: "ethereum" });
+  const eth_account = eth_keyring.createFromUri(
+    `${keys.mnemonic}/m/44'/60'/0'/0/0`,
+  );
+
+  keys.eth_account = {
+    address: eth_account.address,
+    publicKey: u8aToHex(eth_account.publicKey),
+  };
+
+  return keys;
+}
+
+export function getNodeKey(node: Node): GenesisNodeKey {
+  try {
+    const { sr_account, eth_account } = node.accounts;
+
+    const key: GenesisNodeKey = [
+      eth_account.address,
+      eth_account.address,
+      {
+        aura: sr_account.address,
+      },
+    ];
+
+    return key;
+  } catch (err) {
+    console.error(
+      `\n${decorators.red(`Fail to generate key for node: ${node}`)}`,
+    );
+    throw err;
+  }
+}
+
+export async function addCollatorSelection(specPath: string, node: Node) {
+  try {
+    const chainSpec = readAndParseChainSpec(specPath);
+    const runtimeConfig = getRuntimeConfig(chainSpec);
+    if (
+      !runtimeConfig?.collatorSelection?.invulnerables &&
+      !runtimeConfig?.collatorStaking?.invulnerables
+    )
+      return;
+
+    const { eth_account } = node.accounts;
+
+    if (runtimeConfig.collatorSelection)
+      runtimeConfig.collatorSelection.invulnerables.push(eth_account.address);
+    if (runtimeConfig.collatorStaking)
+      runtimeConfig.collatorStaking.invulnerables.push(eth_account.address);
+
+    new CreateLogTable({
+      colWidths: [30, 20, 70],
+    }).pushToPrint([
+      [
+        decorators.cyan("👤 Added CollatorSelection "),
+        decorators.green(node.name),
+        decorators.magenta(eth_account.address),
+      ],
+    ]);
+
+    writeChainSpec(specPath, chainSpec);
+  } catch (err) {
+    console.error(`\n${decorators.red(`Fail to add collator: ${node}`)}`);
+    throw err;
+  }
+}
+
+export default {
+  getNodeKey,
+  generateKeyForNode,
+  addCollatorSelection,
+};
diff --git a/javascript/packages/orchestrator/src/chain-decorators/index.ts b/javascript/packages/orchestrator/src/chain-decorators/index.ts
index 1cd6978b1..2ad981204 100644
--- a/javascript/packages/orchestrator/src/chain-decorators/index.ts
+++ b/javascript/packages/orchestrator/src/chain-decorators/index.ts
@@ -9,6 +9,7 @@ enum CHAIN {
   Oak = "oak",
   Mangata = "mangata",
   Generic = "generic",
+  GenericEvm = "generic_evm",
   LocalV = "local_v",
   MainnetLocalV = "mainnet_local_v",
 }
@@ -29,8 +30,10 @@ import mainnet_local_v from "./mainnet-local-v";
 import mangata from "./mangata";
 import moonbeam from "./moonbeam";
 import oak from "./oak";
+import generic_evm from "./generic-evm";
 
-function whichChain(chain: string): CHAIN {
+function whichChain(chain_name: string, force_decorator?: string): CHAIN {
+  const chain = force_decorator ? force_decorator : chain_name;
   if (chain.includes("statemint") || chain.includes("asset-hub-polkadot"))
     return CHAIN.AssetHubPolkadot;
   if (/moonbase|moonriver|moonbeam/.test(chain)) return CHAIN.Moonbeam;
@@ -43,6 +46,7 @@ function whichChain(chain: string): CHAIN {
   if (/mangata/.test(chain)) return CHAIN.Mangata;
   if (/local-v/.test(chain)) return CHAIN.LocalV;
   if (/mainnet-local-v/.test(chain)) return CHAIN.MainnetLocalV;
+  if (/generic-evm/.test(chain)) return CHAIN.GenericEvm;
 
   return CHAIN.Generic;
 }
@@ -110,6 +114,14 @@ const MainnetLocalVDecorators: Decorator = Object.keys(mainnet_local_v).reduce(
   Object.create({}),
 );
 
+const GenericEvmDecorators: Decorator = Object.keys(generic_evm).reduce(
+  (memo, fn) => {
+    memo[fn] = (generic_evm as Decorator)[fn];
+    return memo;
+  },
+  Object.create({}),
+);
+
 const decorators: { [para in CHAIN]: { [fn: string]: Function } } = {
   moonbeam: moonbeamDecorators,
   asset_hub_polkadot: assetHubPolkadotDecorators,
@@ -123,6 +135,7 @@ const decorators: { [para in CHAIN]: { [fn: string]: Function } } = {
   local_v: localVDecorators,
   mainnet_local_v: MainnetLocalVDecorators,
   generic: {},
+  generic_evm: GenericEvmDecorators,
 };
 
 function decorate(chain: CHAIN, fns: Function[]) {
diff --git a/javascript/packages/orchestrator/src/configGenerator.ts b/javascript/packages/orchestrator/src/configGenerator.ts
index 0da22c9e2..067ac6a0f 100644
--- a/javascript/packages/orchestrator/src/configGenerator.ts
+++ b/javascript/packages/orchestrator/src/configGenerator.ts
@@ -247,6 +247,7 @@ export async function generateNetworkSpec(
     for (const parachain of config.parachains) {
       const para: CHAIN = whichChain(
         parachain.chain || parachain.chain_spec_path || "",
+        parachain.force_decorator,
       );
 
       let computedStatePath,
diff --git a/javascript/packages/orchestrator/src/sharedTypes.ts b/javascript/packages/orchestrator/src/sharedTypes.ts
index 9c4989c2f..0760a4243 100644
--- a/javascript/packages/orchestrator/src/sharedTypes.ts
+++ b/javascript/packages/orchestrator/src/sharedTypes.ts
@@ -147,6 +147,7 @@ export interface Resources {
 export interface CommonParachainConfig {
   id: number;
   chain?: string;
+  force_decorator?: string;
   genesis?: JSON | ObjectJSON;
   balance?: number;
   delayNetworkSettings?: DelayNetworkSettings;