import ast  # The python library not regexploit.ast
import re
from typing import List, Union

from regexploit.found_regex import FoundRegex


RE_FUNC_TO_FLAGS_POS = {
    "compile": 1,
    "search": 2,
    "match": 2,
    "fullmatch": 2,
    "findall": 2,
    "finditer": 2,
    "split": 3,
    "sub": 4,
    "subn": 4,
}


class PythonNodeVisitor(ast.NodeVisitor):
    """
    Try to extract regular expressions from python code by walking the AST.
    """

    def __init__(self):
        self.patterns: List[FoundRegex] = []

    def maybe_pattern(self, lineno: int, pattern: str):
        """Check if the pattern could possibly have ReDoS: if so, add it."""
        if pattern.count("*") + pattern.count("+") + pattern.count(",}") >= 2:
            # Could have ReDoS
            # Now check if it still looks like a docstring
            if " * * *" in pattern:
                return  # Looks like cron (of course could just be really silly regex)
            if pattern.count("\n") < 5 or "?" in pattern or "\\" in pattern:
                self.patterns.append(FoundRegex(lineno, pattern, 0, False))

    def visit_Constant(self, constant: ast.Constant):
        if isinstance(constant.value, bytes):
            try:
                self.maybe_pattern(constant.lineno, constant.value.decode())
            except UnicodeDecodeError:
                pass  # TODO Parse unicode patterns
        elif isinstance(constant.value, str):
            self.maybe_pattern(constant.lineno, constant.value)

    def visit_Assign(self, node: ast.Assign):
        if (
            len(node.targets) != 1
            or not isinstance(node.targets[0], ast.Name)
            or node.targets[0].id != "__doc__"
        ):
            self.generic_visit(node)

    def visit_body_without_docstring(
        self,
        node: Union[ast.FunctionDef, ast.AsyncFunctionDef, ast.Module, ast.ClassDef],
    ):
        if node.body:
            body = node.body
            if isinstance(body[0], ast.Expr):
                potential_docstring = body[0].value
                if isinstance(potential_docstring, ast.Constant):
                    node.body = node.body[1:]  # Ignore docstring

        self.generic_visit(node)

    def visit_FunctionDef(self, node: ast.FunctionDef):
        self.visit_body_without_docstring(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
        self.visit_body_without_docstring(node)

    def visit_ClassDef(self, node: ast.ClassDef):
        self.visit_body_without_docstring(node)

    def visit_Module(self, node: ast.Module):
        self.visit_body_without_docstring(node)

    def visit_Call(self, node: ast.Call):
        if isinstance(node.func, ast.Attribute):
            attr: ast.Attribute = node.func
            if (
                isinstance(attr.value, ast.Name)
                and attr.value.id == "re"
                and len(node.args)
                and isinstance(node.args[0], ast.Constant)
            ):
                flags = 0
                pattern = node.args[0].value
                flags_pos = RE_FUNC_TO_FLAGS_POS.get(attr.attr)
                if flags_pos is not None:
                    # re.compile, re.sub, re.match etc
                    if len(node.args) == flags_pos + 1:
                        flags = RegexFlagVisitor.get_flags(node.args[flags_pos])
                    else:
                        for kw in node.keywords:
                            if kw.arg == "flags":
                                flags = RegexFlagVisitor.get_flags(kw.value)
                                break
                    if isinstance(pattern, bytes):
                        try:
                            pattern = pattern.decode()
                        except UnicodeDecodeError:
                            return  # TODO unicode
                    if isinstance(pattern, str):
                        self.patterns.append(
                            FoundRegex(node.lineno, pattern, flags, True)
                        )
                        return

        self.generic_visit(node)


class RegexFlagVisitor(ast.NodeVisitor):
    """Guess the flags from the 2nd argument of re.compile("abc", re.X | re.M)"""

    def __init__(self):
        self.flags: int = 0

    def visit_Attribute(self, node: ast.Attribute):
        if isinstance(node.value, ast.Name) and node.value.id == "re":
            if node.attr == node.attr.upper():
                try:
                    self.flags |= getattr(re, node.attr)
                except AttributeError:
                    pass

    @staticmethod
    def get_flags(node) -> int:
        rfv = RegexFlagVisitor()
        rfv.visit(node)
        return rfv.flags
