From f77bdb0dd68fbe9197c6102602ff555408271334 Mon Sep 17 00:00:00 2001
From: Pavel Shuvalov <pavel.s@ton.org>
Date: Fri, 6 Sep 2024 11:04:53 +0400
Subject: [PATCH] Extract list of enrolled users from cNFT

---
 backends/sbt_enrollment.py | 60 +++++++++++++++++++++++++++++++++++---
 models/season_config.py    |  5 ++++
 test_runner.py             |  8 +++++
 3 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/backends/sbt_enrollment.py b/backends/sbt_enrollment.py
index 0006431..0a11b9e 100644
--- a/backends/sbt_enrollment.py
+++ b/backends/sbt_enrollment.py
@@ -1,3 +1,14 @@
+from backends.tonapi import TonapiAdapter
+from backends.contracts_executor import ContractsExecutor
+from models.season_config import SeasonConfig
+from loguru import logger
+import requests
+import psycopg2
+import time
+import base64
+from tonsdk.boc import Cell
+from tonsdk.utils import Address
+
 """
 Simple tool to sync cNFT SBT collection used for enrollment. 
 1. Fetches latest cNFT merkle proof hash
@@ -12,11 +23,52 @@
   sbt varchar,
   added_at timestamp
 )
+
+CREATE UNIQUE INDEX enrollment_{season_name}_address_idx ON tol.enrollment_{season_name (address);
 """
+API_BASE_URL = "https://stg.globalsociety.cc/v1/csbts/"
 
-class SBTEnrollmentSync(CalculationBackend):
-    def __init__(self, connection):
+class SBTEnrollmentSync:
+    def __init__(self, connection, tonapi: TonapiAdapter, executor: ContractsExecutor):
         self.connection = connection
+        self.tonapi = tonapi
+        self.executor = executor
 
-    def sync(self, config: SeasonConfig)
-        # TODO
\ No newline at end of file
+    def sync(self, config: SeasonConfig):
+        logger.info(f"Requesting state for {config.enrollment_sbt}")
+        end_block = None
+        if int(time.time()) > config.end_time:
+            end_block = config.block_before_end_ref
+            assert end_block is not None, "Season is closed, one need to specify last block ref"
+        code, data = self.tonapi.get_state(config.enrollment_sbt, target_block=end_block)
+        # logger.info(f"Got state: {state}")
+        [merkle_root] = self.executor.execute(code, data, config.enrollment_sbt, 'get_merkle_root', ['int'])
+        merkle_root_hex = f"{int(merkle_root):064x}"
+        logger.info(f"Requesting data for {merkle_root_hex}")
+        start = 0
+        PAGE = 100 # increase to 1000
+        total = None
+        owners = set()
+        while True:
+          part = requests.get(f"{API_BASE_URL}{merkle_root_hex}/items?_start={start}&_end={start + PAGE}").json()
+          assert total is None or total == part['data']['total']
+          total = part['data']['total']
+          for item in part['data']['items']:
+            data_cell = item['data_cell']
+            cell = Cell.one_from_boc(base64.b64decode(data_cell))
+            # TODO check hash and extract owner
+            owner = Address(item['metadata']['owner']).to_string(1, 1, 1)
+            owners.add(owner)
+          if len(owners) == total:
+            break
+          else:
+            logger.info(f"Got {len(owners)} items, need {total}")
+            start += PAGE
+        logger.info(f"Got {len(owners)} owners to update")
+        with self.connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
+          for owner in owners:
+            cursor.execute(f"""insert into tol.enrollment_{config.safe_season_name()}(address, added_at)
+            values (%s, now())
+            on conflict do nothing
+            """, (owner, ))
+        self.connection.commit()
diff --git a/models/season_config.py b/models/season_config.py
index fab210c..569d0f0 100644
--- a/models/season_config.py
+++ b/models/season_config.py
@@ -50,6 +50,11 @@ class SeasonConfig:
     block_before_start_ref: str = None # ref for the last block exactly before the start_time
     block_before_end_ref: str = None # ref for the last block of the season (i.e. before end_time)
 
+
+    """
+    Address of SBT collection which is used for enrollment
+    """
+    enrollment_sbt: str = None 
     # for SQL usage
     def safe_season_name(self):
         return self.name.replace(".", '_')
\ No newline at end of file
diff --git a/test_runner.py b/test_runner.py
index 96da146..80f90b7 100644
--- a/test_runner.py
+++ b/test_runner.py
@@ -9,6 +9,7 @@
 
 from backends.contracts_executor import ContractsExecutor
 from backends.defi import DefillamaDeFiBackend
+from backends.sbt_enrollment import SBTEnrollmentSync
 from backends.redoubt.apps import RedoubtAppBackend
 from backends.redoubt.apps_v2 import RedoubtAppBackendV2
 from backends.redoubt.nfts import RedoubtNFTsBackend
@@ -41,6 +42,13 @@
                 executor=ContractsExecutor(os.getenv('CONTRACTS_EXECUTOR_URL'))
             )
             season = S5_defi
+        elif sys.argv[1] == 'sbt':
+            backend = SBTEnrollmentSync(conn,
+                tonapi=TonapiAdapter(),
+                executor=ContractsExecutor(os.getenv('CONTRACTS_EXECUTOR_URL')))
+            season = S6_apps
+            backend.sync(season)
+            sys.exit(0)
         else:
             raise Exception(f"leaderboard not supported: {sys.argv[1]}")