import fs from 'fs' ; import path , { join } from 'path' ; import os from 'os' ; import { PythonShell } from 'python-shell' ; import ReadLines from 'n-readlines' ; import { BlockEvent , Finding , HandleBlock , HandleTransaction , TransactionEvent } from "../../sdk" ; export type GetPythonAgentHandlers = ( pythonAgentPath : string ) => Promise < { handleTransaction ? : HandleTransaction ; handleBlock ? : HandleBlock ; } > ; const INITIALIZE_MARKER = "!*initialize*!" ; const FINDING_MARKER = "!*forta_finding*!:" ; const INITIALIZE_METHOD_NAME = "initialize" ; const HANDLE_TRANSACTION_METHOD_NAME = 'handle_transaction' ; const HANDLE_BLOCK_METHOD_NAME = 'handle_block' ; export function provideGetPythonAgentHandlers ( ) { return async function getPythonAgentHandlers ( pythonAgentPath ) { const { hasInitializeHandler , hasBlockHandler , hasTransactionHandler } = hasHandlers ( pythonAgentPath ) ; if ( ! hasBlockHandler && ! hasTransactionHandler ) throw new Error ( `no handlers found in ${ pythonAgentPath }` ) ; const pythonHandler = getPythonHandler ( pythonAgentPath ) ; return { initialize : hasInitializeHandler ? pythonHandler : undefined , handleBlock : hasBlockHandler ? pythonHandler : undefined , handleTransaction : hasTransactionHandler ? pythonHandler : undefined } ; } ; } function hasHandlers ( agentPath ) { const lineReader = new ReadLines ( agentPath ) ; let hasTransactionHandler = false ; let hasBlockHandler = false ; let hasInitializeHandler = false ; let line ; while ( line = lineReader . next ( ) ) { line = line . toString ( 'ascii' ) ; if ( line . startsWith ( `def ${ HANDLE_TRANSACTION_METHOD_NAME }` ) ) { hasTransactionHandler = true ; } else if ( line . startsWith ( `def ${ HANDLE_BLOCK_METHOD_NAME }` ) ) { hasBlockHandler = true ; } else if ( line . startsWith ( `def ${ INITIALIZE_METHOD_NAME }` ) ) { hasInitializeHandler = true ; } } return { hasTransactionHandler , hasBlockHandler , hasInitializeHandler } ; } function getPythonHandler ( agentPath ) { const agentModule = agentPath . replace ( `${ process . cwd ( ) }${ path . sep }` , '' ) . replace ( '.py' , '' ) . replace ( path . sep , '.' ) ; const pythonWrapperScript = `
import sys
sys.path.append('${ process . cwd ( ) }')
import json
from forta_agent import BlockEvent, TransactionEvent
import ${ agentModule }

while True:
  try:
    msgJson = json.loads(input())
    msgType = msgJson['msgType']
    if msgType == '${ INITIALIZE_METHOD_NAME }':
      ${ agentModule }.${ INITIALIZE_METHOD_NAME }()
      print(f'${ INITIALIZE_MARKER }')
    elif msgType == '${ HANDLE_TRANSACTION_METHOD_NAME }':
      hash = msgJson['hash']
      event = TransactionEvent(msgJson)
      findings = ${ agentModule }.${ HANDLE_TRANSACTION_METHOD_NAME }(event)
      print(f'${ FINDING_MARKER }{hash}:{json.dumps(findings, default=lambda f: f.toJson())}')
    elif msgType == '${ HANDLE_BLOCK_METHOD_NAME }':
      hash = msgJson['hash']
      event = BlockEvent(msgJson)
      findings = ${ agentModule }.${ HANDLE_BLOCK_METHOD_NAME }(event)
      print(f'${ FINDING_MARKER }{hash}:{json.dumps(findings, default=lambda f: f.toJson())}')
  except Exception as e:
    print(e, file=sys.stderr)
` ; const randomInt = Math . floor ( Math . random ( ) * Date . now ( ) ) ; const pythonWrapperPath = join ( os . tmpdir ( ) , `forta_agent_${ randomInt }.py` ) ; fs . writeFileSync ( pythonWrapperPath , pythonWrapperScript ) ; PythonShell . checkSyntaxFile ( pythonWrapperPath ) ; const promiseCallbackMap = { } ; const pythonWrapper = new PythonShell ( pythonWrapperPath , { args : process . argv } ) . on ( "message" , function ( message ) { if ( message . startsWith ( INITIALIZE_MARKER ) ) { const { resolve } = promiseCallbackMap [ 'init' ] ; resolve ( undefined ) ; delete promiseCallbackMap [ 'init' ] ; return ; } else if ( ! message . startsWith ( FINDING_MARKER ) ) { console . log ( message ) ; return ; } const hashStartIndex = FINDING_MARKER . length ; const hashEndIndex = message . indexOf ( ':' , hashStartIndex ) ; const hash = message . substr ( hashStartIndex , hashEndIndex - hashStartIndex ) ; const findingsJson = JSON . parse ( message . substr ( hashEndIndex + 1 ) ) as any [ ] ; const findings = findingsJson . map ( findingJson => Finding . fromObject ( JSON . parse ( findingJson ) ) ) ; const { resolve } = promiseCallbackMap [ hash ] ; resolve ( findings ) ; delete promiseCallbackMap [ hash ] ; } ) . on ( "stderr" , function ( err ) { console . log ( err ) ; const hash = Object . keys ( promiseCallbackMap ) [ 0 ] ; if ( hash && promiseCallbackMap [ hash ] ) { const { reject } = promiseCallbackMap [ hash ] ; reject ( new Error ( `python: ${ err }` ) ) ; } } ) ; return function handler ( event ) { return new Promise ( ( resolve , reject ) => { let msgType , hash ; if ( ! event ) { msgType = INITIALIZE_METHOD_NAME ; hash = 'init' ; } else if ( event instanceof BlockEvent ) { msgType = HANDLE_BLOCK_METHOD_NAME ; hash = event . blockHash ; } else { msgType = HANDLE_TRANSACTION_METHOD_NAME ; hash = event . hash ; } promiseCallbackMap [ hash ] = { resolve , reject } ; pythonWrapper . send ( JSON . stringify ( { msgType , hash , ... event } ) ) ; } ) ; } as any ; }