import re

from core.colors import end, red, yellow

if len(end) < 1:
    end = red = yellow = "*"


def dom(response):
    highlighted = []
    sources = r"""\b(?:document\.(URL|documentURI|URLUnencoded|baseURI|cookie|referrer)|location\.(href|search|hash|pathname)|window\.name|history\.(pushState|replaceState)(local|session)Storage)\b"""
    sinks = r"""\b(?:eval|evaluate|execCommand|assign|navigate|getResponseHeaderopen|showModalDialog|Function|set(Timeout|Interval|Immediate)|execScript|crypto.generateCRMFRequest|ScriptElement\.(src|text|textContent|innerText)|.*?\.onEventName|document\.(write|writeln)|.*?\.innerHTML|Range\.createContextualFragment|(document|window)\.location)\b"""
    scripts = re.findall(r"(?i)(?s)<script[^>]*>(.*?)</script>", response)
    sinkFound, sourceFound = False, False
    for script in scripts:
        script = script.split("\n")
        num = 1
        allControlledVariables = set()
        try:
            for newLine in script:
                line = newLine
                parts = line.split("var ")
                controlledVariables = set()
                if len(parts) > 1:
                    for part in parts:
                        for controlledVariable in allControlledVariables:
                            if controlledVariable in part:
                                controlledVariables.add(
                                    re.search(r"[a-zA-Z$_][a-zA-Z0-9$_]+", part)
                                    .group()
                                    .replace("$", "\$")
                                )
                pattern = re.finditer(sources, newLine)
                for grp in pattern:
                    if grp:
                        source = newLine[grp.start() : grp.end()].replace(" ", "")
                        if source:
                            if len(parts) > 1:
                                for part in parts:
                                    if source in part:
                                        controlledVariables.add(
                                            re.search(r"[a-zA-Z$_][a-zA-Z0-9$_]+", part)
                                            .group()
                                            .replace("$", "\$")
                                        )
                            line = line.replace(source, yellow + source + end)
                for controlledVariable in controlledVariables:
                    allControlledVariables.add(controlledVariable)
                for controlledVariable in allControlledVariables:
                    matches = list(
                        filter(None, re.findall(r"\b%s\b" % controlledVariable, line))
                    )
                    if matches:
                        sourceFound = True
                        line = re.sub(
                            r"\b%s\b" % controlledVariable,
                            yellow + controlledVariable + end,
                            line,
                        )
                pattern = re.finditer(sinks, newLine)
                for grp in pattern:
                    if grp:
                        sink = newLine[grp.start() : grp.end()].replace(" ", "")
                        if sink:
                            line = line.replace(sink, red + sink + end)
                            sinkFound = True
                if line != newLine:
                    highlighted.append("%-3s %s" % (str(num), line.lstrip(" ")))
                num += 1
        except MemoryError:
            pass
    if sinkFound or sourceFound:
        return highlighted
    else:
        return []
