Coverage for webapp/extensions.py: 35%

1import base64 

2import json 

3import posixpath 

4import re 

5from os import path 

6from pathlib import Path 

7from urllib.parse import urljoin 

8 

9import flask 

10from canonicalwebteam.flask_vite import FlaskVite 

11from canonicalwebteam.flask_vite.extension import ( 

12 vite_import as flask_vite_import, 

13) 

14from flask_wtf.csrf import CSRFProtect 

15from markupsafe import Markup 

16 

17csrf = CSRFProtect() 

18vite = FlaskVite() 

19 

20_SCRIPT_EXTENSIONS = {".js", ".ts", ".jsx", ".tsx", ".svelte", ".vue"} 

21_IMPORT_SPECIFIER_PATTERN = re.compile( 

22 r'(?P<prefix>\bfrom\s*["\']|\bimport\s*["\']|\bimport\s*\(\s*["\'])' 

23 r'(?P<specifier>[^"\']+)' 

24 r'(?P<suffix>["\'])' 

25) 

26_IMPORT_META_URL_PATTERN = re.compile(r"\bimport\.meta\.url\b") 

27 

28 

29def _csp_nonce_attr() -> str: 

30 nonce = getattr(flask.request, "CSP_NONCE", "") 

31 if not nonce: 

32 return "" 

33 return f' nonce="{nonce}"' 

34 

35 

36def _escape_inline_script(content: str) -> str: 

37 return content.replace("</script", "<\\/script") 

38 

39 

40def _resolve_specifier(module_url: str, specifier: str) -> str: 

41 if specifier.startswith("/"): 

42 return posixpath.normpath(specifier) 

43 

44 if specifier.startswith("."): 

45 module_dir = posixpath.dirname(module_url) 

46 return posixpath.normpath(posixpath.join(module_dir, specifier)) 

47 

48 return specifier 

49 

50 

51def _is_local_vite_script(specifier: str) -> bool: 

52 _, extension = path.splitext(specifier) 

53 return ( 

54 specifier.startswith("/static/js/dist/vite/") 

55 and extension in _SCRIPT_EXTENSIONS 

56 ) 

57 

58 

59def _extract_local_vite_imports( 

60 module_source: str, module_url: str 

61) -> list[str]: 

62 discovered_imports: list[str] = [] 

63 

64 for match in _IMPORT_SPECIFIER_PATTERN.finditer(module_source): 

65 specifier = match.group("specifier") 

66 resolved_specifier = _resolve_specifier(module_url, specifier) 

67 

68 if _is_local_vite_script(resolved_specifier): 

69 discovered_imports.append(resolved_specifier) 

70 

71 return discovered_imports 

72 

73 

74def _rewrite_module_specifiers( 

75 module_source: str, module_url: str, known_module_urls: set[str] 

76) -> str: 

77 def to_absolute_url(specifier_path: str) -> str: 

78 return urljoin(flask.request.url_root, specifier_path.lstrip("/")) 

79 

80 def replace_specifier(match: re.Match) -> str: 

81 specifier = match.group("specifier") 

82 resolved_specifier = _resolve_specifier(module_url, specifier) 

83 

84 if resolved_specifier in known_module_urls: 

85 absolute_url = to_absolute_url(resolved_specifier) 

86 return ( 

87 f"{match.group('prefix')}{absolute_url}" 

88 f"{match.group('suffix')}" 

89 ) 

90 

91 if specifier.startswith("/") or specifier.startswith("."): 

92 absolute_url = to_absolute_url(resolved_specifier) 

93 return ( 

94 f"{match.group('prefix')}{absolute_url}" 

95 f"{match.group('suffix')}" 

96 ) 

97 

98 return match.group(0) 

99 

100 rewritten_source = _IMPORT_SPECIFIER_PATTERN.sub( 

101 replace_specifier, module_source 

102 ) 

103 module_absolute_url = urljoin( 

104 flask.request.url_root, module_url.lstrip("/") 

105 ) 

106 return _IMPORT_META_URL_PATTERN.sub( 

107 json.dumps(module_absolute_url), rewritten_source 

108 ) 

109 

110 

111def _read_asset(url: str) -> str: 

112 file_path = url.lstrip("/") 

113 with open(file_path, encoding="utf-8") as asset: 

114 return asset.read() 

115 

116 

117def _get_vite_script_module_urls() -> list[str]: 

118 out_dir = Path( 

119 flask.current_app.config.get("VITE_OUTDIR", "static/js/dist/vite") 

120 ) 

121 module_urls = [ 

122 f"/{module_path.as_posix()}" 

123 for module_path in out_dir.rglob("*.js") 

124 if module_path.is_file() 

125 ] 

126 module_urls.sort() 

127 return module_urls 

128 

129 

130def _get_rewritten_inline_modules() -> dict[str, str]: 

131 cached_modules = getattr(flask.g, "_vite_inline_rewritten_modules", None) 

132 if cached_modules is not None: 

133 return cached_modules 

134 

135 module_urls = _get_vite_script_module_urls() 

136 known_module_urls = set(module_urls) 

137 rewritten_modules = { 

138 module_url: _rewrite_module_specifiers( 

139 _read_asset(module_url), module_url, known_module_urls 

140 ) 

141 for module_url in module_urls 

142 } 

143 

144 flask.g._vite_inline_rewritten_modules = rewritten_modules 

145 return rewritten_modules 

146 

147 

148def _build_import_map_script( 

149 rewritten_modules: dict[str, str], nonce_attr: str 

150) -> str: 

151 import_map = { 

152 "imports": { 

153 urljoin(flask.request.url_root, module_url.lstrip("/")): ( 

154 "data:text/javascript;base64," 

155 + base64.b64encode(module_code.encode("utf-8")).decode("utf-8") 

156 ) 

157 for module_url, module_code in rewritten_modules.items() 

158 } 

159 } 

160 import_map_json = json.dumps(import_map, separators=(",", ":")) 

161 return ( 

162 f'<script type="importmap"{nonce_attr}>' 

163 f"{_escape_inline_script(import_map_json)}" 

164 "</script>" 

165 ) 

166 

167 

168def _inline_script_import(entrypoint: str) -> Markup: 

169 entry_url = vite.instance.get_asset_url(entrypoint) 

170 css_urls = vite.instance.get_imported_css(entrypoint) 

171 rewritten_modules = _get_rewritten_inline_modules() 

172 

173 nonce_attr = _csp_nonce_attr() 

174 import_map_script = "" 

175 if not getattr(flask.g, "_vite_inline_import_map_emitted", False): 

176 import_map_script = _build_import_map_script( 

177 rewritten_modules, nonce_attr 

178 ) 

179 flask.g._vite_inline_import_map_emitted = True 

180 

181 entry_script = ( 

182 f'<script type="module"{nonce_attr}>' 

183 f"{_escape_inline_script(rewritten_modules[entry_url])}" 

184 "</script>" 

185 ) 

186 css_links = "".join( 

187 f'<link rel="stylesheet" href="{css_url}"{nonce_attr} />' 

188 for css_url in css_urls 

189 ) 

190 

191 return Markup("".join([import_map_script, entry_script, css_links])) 

192 

193 

194def vite_import(entrypoint: str) -> Markup: 

195 _, extension = path.splitext(entrypoint) 

196 is_inline_script = extension in _SCRIPT_EXTENSIONS 

197 is_prod = flask.current_app.config.get("VITE_MODE") == "production" 

198 inline_enabled = flask.current_app.config.get("VITE_INLINE_JS", False) 

199 

200 if not (is_inline_script and is_prod and inline_enabled): 

201 return flask_vite_import(entrypoint) 

202 

203 return _inline_script_import(entrypoint)