Module refinery.units.scripting

from __future__ import annotations

import abc
import codecs

from refinery.lib.id import guess_text_encoding
from refinery.lib.scripts import TREE_RECURSION_DEPTH, Node
from refinery.lib.scripts.pipeline import DeobfuscationTimeout
from refinery.lib.tools import RecursionDepth
from refinery.lib.types import Param, buf
from refinery.units import Arg, Chunk, RefineryPartialResult, Unit


class AutoDeobfuscationTimeout(RefineryPartialResult):
    def __init__(self, partial, reason: str | None = None):
        message = 'The deobfuscation timeout was reached before the data stabilized'
        message = F'{message}; {reason}.' if reason else F'{message}.'
        super().__init__(message, partial=partial)


class IterativeDeobfuscator(Unit, abstract=True):

    def __init__(
        self,
        timeout: Param[int, Arg.Number(
            '-t', help='Maximum number of iterations; the default is {default}.')] = 500,
        **kwargs,
    ):
        super().__init__(timeout=timeout, **kwargs)

    @abc.abstractmethod
    def parse(self, data: str) -> Node:
        ...

    @abc.abstractmethod
    def transform(self, ast: Node) -> int:
        ...

    @abc.abstractmethod
    def synthesize(self, ast: Node) -> str:
        ...

    def process(self, data: Chunk) -> buf:
        try:
            with RecursionDepth(TREE_RECURSION_DEPTH):
                return self._process(data)
        except RecursionError:
            self.log_warn('input nesting exceeded the interpreter stack; leaving it unchanged')
            return data

    def _process(self, data: Chunk) -> buf:
        self.log_info('parsing input data')
        view = memoryview(data)

        if format := guess_text_encoding(data):
            codec = format.codec
            view = view[format.bom:]
            errors = 'surrogatepass' if format.step > 1 else 'surrogateescape'
        else:
            codec = self.codec
            errors = 'surrogateescape'

        def _result():
            return codecs.encode(
                self.synthesize(ast), self.codec, errors=errors)

        txt = codecs.decode(view, codec, errors=errors)
        ast = self.parse(txt)

        for k in range(self.args.timeout):
            self.log_info(F'starting round {k}')
            try:
                changed = self.transform(ast)
            except DeobfuscationTimeout as error:
                raise AutoDeobfuscationTimeout(_result(), reason=str(error)) from error
            if changed:
                continue
            return _result()
        else:
            raise AutoDeobfuscationTimeout(_result())