userscript-proxy/injector.py at develop · lydell/userscript-proxy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
from typing import Optional, Iterable, List, Callable, Pattern, Match, Tuple
import glob, os
from bs4 import BeautifulSoup, Comment, Doctype
from mitmproxy import ctx, http
from functools import partial
import shlex
import warnings
from modules.metadata import MetadataError, PREFIX_TAG
import modules.userscript as userscript
import modules.inline as inline
import modules.text as T
from modules.userscript import Userscript, UserscriptError, document_end, document_start, document_idle
from modules.utilities import first, second, itemList, fromOptional, flag, idem
from modules.constants import VERSION, VERSION_PREFIX, APP_NAME, DEFAULT_USERSCRIPTS_DIR, DEFAULT_QUERY_PARAM_TO_DISABLE
from modules.inject import Options, inject
from modules.misc import sanitize
from modules.requests import CONTENT_TYPE, inferEncoding, requestContainsQueryParam

PATTERN_USERSCRIPT: str = "*.user.js"
RELEVANT_CONTENT_TYPES: List[str] = ["text/html", "application/xhtml+xml"]
CHARSET_DEFAULT: str = "utf-8"
TAB: str = "    "
LIST_ITEM_PREFIX: str = TAB + "• "
HTML_PARSER: str = "lxml"
# lxml handles non-uppercase DOCTYPE correctly; html.parser does not: It emits
# <!DOCTYPE doctype html> if the original source code contained <!doctype html>.
HTML_INFO_COMMENT_PREFIX: str = f"""
[{T.INFO_MESSAGE}]
"""


def logInfo(s: str) -> None:
    try:
        ctx.log.info(s)
    except Exception:
        print(s)

def logWarning(s: str) -> None:
    try:
        ctx.log.warn(s)
    except Exception:
        print(s)

def logError(s: str) -> None:
    try:
        ctx.log.error(s)
    except Exception:
        print(s)

def indexOfDTD(soup: BeautifulSoup) -> Optional[int]:
    index: int = 0
    for item in soup.contents:
        if isinstance(item, Doctype):
            return index
        index += 1
    return None

bulletList: Callable[[Iterable[str]], str] = partial(itemList, LIST_ITEM_PREFIX)

def unsafeSequencesMessage(script: Userscript) -> str:
    sequences = script.unsafeSequences
    return f"""{script.name} cannot be injected because it contains {"these unsafe sequences" if len(sequences) > 1 else "this unsafe sequence"}:

{itemList(TAB, sequences)}

<script> tags cannot contain any of these sequences (case-insensitive):

{itemList(TAB, inline.DANGEROUS_SEQUENCES)}

Possible solutions:
""" + bulletList([
    f"Make sure the userscript does not contain any of the sequences listed above.",
    f"Make the userscript available online and give it a {PREFIX_TAG}{userscript.directive_downloadURL}",
    f"Remove the {flag(T.option_inline)} flag.",
])


# Because ctx.options is not subscriptable and we want to be able to use
# expressions as keys:
def option(key: str):
    return ctx.options.__getattr__(sanitize(key))


def loadUserscripts(directories: Iterable[str], recursive: bool) -> List[Userscript]:
    logInfo("Loading userscripts ...")
    loadedUserscripts: List[Tuple[Userscript, str]] = []
    workingDirectory = os.getcwd()
    for directory in directories:
        logInfo(f"""Looking{" recursively" if recursive else ""} for userscripts ({PATTERN_USERSCRIPT}) in directory `{directory}` ...""")
        if not recursive:
            logInfo(f"{TAB}(use {flag(T.option_recursive)} to look recursively)")
        try:
            os.chdir(directory)
        except FileNotFoundError:
            logWarning("Directory `"+directory+"` does not exist.")
            continue
        except PermissionError:
            logError("Permission was denied when trying to read directory `"+directory+"`.")
            continue
        pattern = ("**/" if recursive else "") + PATTERN_USERSCRIPT
        # recursive=True only affects the meaning of "**".
        # https://docs.python.org/3/library/glob.html#glob.glob
        for unsafe_filename in glob.glob(pattern, recursive=True):
            filename = shlex.quote(unsafe_filename)
            logInfo("Loading " + filename + " ...")
            try:
                content = open(filename).read()
            except PermissionError:
                logError("Could not read file `"+filename+"`: Permission denied.")
                continue
            except Exception as e:
                logError("Could not read file `"+filename+"`: " + str(e))
                continue
            try:
                script = userscript.create(content)
                loadedUserscripts.append((script, filename))
                if script.downloadURL is None and len(script.unsafeSequences) > 0:
                    logError(unsafeSequencesMessage(script))
            except MetadataError as err:
                logError("Metadata error:")
                logError(str(err))
                continue
            except UserscriptError as err:
                logError("Userscript error:")
                logError(str(err))
                continue
        os.chdir(workingDirectory) # so mitmproxy does not unload the script
    logInfo("")
    logInfo(str(len(loadedUserscripts)) + " userscript(s) loaded:")
    logInfo(bulletList(map(
        lambda s: f"{first(s).name} ({second(s)})",
        loadedUserscripts
    )))
    logInfo("")
    return list(map(first, loadedUserscripts))


class UserscriptInjector:
    def __init__(self):
        self.userscripts: List[Userscript] = []


    def load(self, loader):
        loader.add_option(sanitize(T.option_inline), bool, False, T.help_inline)
        loader.add_option(sanitize(T.option_recursive), bool, False, T.help_recursive)
        loader.add_option(sanitize(T.option_list_injected), bool, False, T.help_list_injected)
        loader.add_option(sanitize(T.option_userscripts), str, DEFAULT_USERSCRIPTS_DIR, T.help_userscripts)
        loader.add_option(sanitize(T.option_query_param_to_disable), str, DEFAULT_QUERY_PARAM_TO_DISABLE, T.help_query_param_to_disable)


    def configure(self, updates):
        if sanitize(T.option_inline) in updates and option(T.option_inline):
            logWarning(f"""Only inline injection will be used due to {flag(T.option_inline)} flag.""")
        if sanitize(T.option_query_param_to_disable) in updates:
            logInfo(f"""Userscripts will not be injected when the request URL contains a `{option(T.option_query_param_to_disable)}` query parameter.""")
        if sanitize(T.option_userscripts) in updates:
            self.userscripts = loadUserscripts(
                directories = [ option(T.option_userscripts) ],
                recursive = option(T.option_recursive),
            )


    def response(self, flow: http.HTTPFlow):
        response = flow.response
        if CONTENT_TYPE in response.headers:
            if any(map(lambda t: t in response.headers[CONTENT_TYPE], RELEVANT_CONTENT_TYPES)):
                # Response is a web page; proceed.
                insertedScripts: List[str] = []
                soup = BeautifulSoup(
                    response.content,
                    HTML_PARSER,
                    from_encoding=inferEncoding(response)
                )
                requestURL = flow.request.pretty_url # should work in transparent mode too, unless the Host header is spoofed
                if requestContainsQueryParam(option(T.option_query_param_to_disable), flow.request):
                    logInfo(f"""Not injecting any userscripts into {requestURL} because it contains a `{option(T.option_query_param_to_disable)}` query parameter.""")
                    return
                isApplicable: Callable[[Userscript], bool] = userscript.applicableChecker(requestURL)
                for script in self.userscripts:
                    if isApplicable(script):
                        useInline = option(T.option_inline) or script.downloadURL is None
                        if useInline and len(script.unsafeSequences) > 0:
                            logError(unsafeSequencesMessage(script))
                            continue
                        logInfo(f"""Injecting {script.name}{"" if script.version is None else " " + VERSION_PREFIX + script.version} into {requestURL} ({"inline" if useInline else "linked"}) ...""")
                        result = inject(script, soup, Options(
                            inline = option(T.option_inline),
                        ))
                        if type(result) is BeautifulSoup:
                            soup = result
                            insertedScripts.append(script.name + ("" if script.version is None else " " + T.stringifyVersion(script.version)))
                        else:
                            logError("Injection failed due to the following error:")
                            logError(str(result))

                index_DTD: Optional[int] = indexOfDTD(soup)
                # Insert information comment:
                if option(T.option_list_injected):
                    soup.insert(0 if index_DTD is None else 1+index_DTD, Comment(
                        HTML_INFO_COMMENT_PREFIX + (
                            "No matching userscripts for this URL." if insertedScripts == []
                            else "These scripts were inserted:\n" + bulletList(insertedScripts)
                        ) + "\n"
                    ))
                # Serialize and encode:
                response.content = str(soup).encode(
                    fromOptional(soup.original_encoding, CHARSET_DEFAULT),
                    "replace"
                )


addons = [ UserscriptInjector() ]