Coverage for an_website / update / update.py: 88.889%

1# This program is free software: you can redistribute it and/or modify 

2# it under the terms of the GNU Affero General Public License as 

3# published by the Free Software Foundation, either version 3 of the 

4# License, or (at your option) any later version. 

5# 

6# This program is distributed in the hope that it will be useful, 

7# but WITHOUT ANY WARRANTY; without even the implied warranty of 

8# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

9# GNU Affero General Public License for more details. 

10# 

11# You should have received a copy of the GNU Affero General Public License 

12# along with this program. If not, see <https://www.gnu.org/licenses/>. 

13 

14"""The API for updating the website.""" 

15 

16import asyncio 

17import logging 

18import os 

19import sys 

20from asyncio import Future 

21from queue import SimpleQueue 

22from tempfile import NamedTemporaryFile, TemporaryDirectory 

23from typing import Any, ClassVar, Final, Protocol 

24from urllib.parse import unquote 

25 

26from tornado.web import stream_request_body 

27 

28from .. import EVENT_SHUTDOWN, NAME 

29from ..utils.decorators import requires 

30from ..utils.request_handler import APIRequestHandler 

31from ..utils.utils import ModuleInfo, Permission 

32 

33LOGGER: Final = logging.getLogger(__name__) 

34 

35 

36class TempFile(Protocol): 

37 """Should inherit IO[bytes], but that isn't a protocol.""" 

38 

39 def close(self) -> None: 

40 """Close the file.""" 

41 

42 def flush(self) -> None: 

43 """Flush the data down the skibidi toilet.""" 

44 

45 @property 

46 def name(self) -> str: 

47 """The name of the file.""" 

48 

49 def write(self, data: bytes, /) -> int: 

50 """Write data to the file.""" 

51 

52 

53def get_module_info() -> ModuleInfo: 

54 """Create and return the ModuleInfo for this module.""" 

55 return ModuleInfo( 

56 handlers=((r"/api/update/(.*)", UpdateAPI),), 

57 name="Update-API", 

58 description=f"API zum Aktualisieren von {NAME.removesuffix('-dev')}", 

59 hidden=True, 

60 ) 

61 

62 

63def write_from_queue(file: TempFile, queue: SimpleQueue[None | bytes]) -> None: 

64 """Read from a queue and write to a file.""" 

65 while (chunk := queue.get()) is not None: # pylint: disable=while-used 

66 file.write(chunk) 

67 file.flush() 

68 

69 

70@stream_request_body 

71class UpdateAPI(APIRequestHandler): # pragma: no cover 

72 """The request handler for the update API.""" 

73 

74 ALLOWED_METHODS: ClassVar[tuple[str, ...]] = ("PUT",) 

75 POSSIBLE_CONTENT_TYPES: ClassVar[tuple[str, ...]] = ("text/plain",) 

76 

77 dir: TemporaryDirectory[str] 

78 file: TempFile 

79 queue: SimpleQueue[None | bytes] 

80 future: Future[Any] 

81 

82 def data_received(self, chunk: bytes) -> None: # noqa: D102 

83 self.queue.put(chunk) 

84 

85 def on_finish(self) -> None: # noqa: D102 

86 if hasattr(self, "queue"): 

87 self.queue.put(None) 

88 

89 async def pip_install(self, *args: str) -> int: 

90 """Install something and write the output.""" 

91 process = await asyncio.create_subprocess_exec( 

92 sys.executable, 

93 "-m", 

94 "pip", 

95 "install", 

96 "--require-virtualenv", 

97 *args, 

98 stdin=asyncio.subprocess.DEVNULL, 

99 stdout=asyncio.subprocess.PIPE, 

100 stderr=asyncio.subprocess.STDOUT, 

101 ) 

102 # pylint: disable=while-used 

103 while not process.stdout.at_eof(): # type: ignore[union-attr] 

104 char = await process.stdout.read(1) # type: ignore[union-attr] 

105 self.write(char) 

106 if char == b"\n": 

107 self.flush() # type: ignore[unused-awaitable] 

108 await process.wait() 

109 assert process.returncode is not None 

110 return process.returncode 

111 

112 async def prepare(self) -> None: # noqa: D102 

113 await super().prepare() 

114 loop = asyncio.get_running_loop() 

115 self.dir = TemporaryDirectory() 

116 self.file = NamedTemporaryFile(dir=self.dir.name, delete=False) 

117 self.queue = SimpleQueue() 

118 self.future = loop.run_in_executor( 

119 None, write_from_queue, self.file, self.queue 

120 ) 

121 

122 @requires(Permission.UPDATE) 

123 async def put(self, filename: str) -> None: 

124 """Handle PUT requests to the update API.""" 

125 self.queue.put(None) 

126 await self.future 

127 self.file.close() 

128 

129 filepath = os.path.join(self.dir.name, unquote(filename)) 

130 os.rename(self.file.name, filepath) 

131 

132 self.set_status(202) 

133 self.set_header("X-Accel-Buffering", "no") 

134 

135 await self.pip_install("--upgrade", "pip") 

136 

137 returncode = await self.pip_install(filepath) 

138 

139 await self.finish() 

140 

141 if returncode: 

142 LOGGER.error("Failed to install %s", filename) 

143 elif self.get_bool_argument("shutdown", True): 

144 EVENT_SHUTDOWN.set()