diff --git a/CHANGELOG.md b/CHANGELOG.md index 7adfa30..e42a1fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Next version +- add cookie support - add !restart to botrulez - save (overwrite) `Bot` config file diff --git a/README.md b/README.md index 5e2e933..94bef94 100644 --- a/README.md +++ b/README.md @@ -59,7 +59,6 @@ i. e. the text between the end of the "!echo" and the end of the whole message. ## TODOs - [ ] document yaboli (markdown files in a "docs" folder?) -- [ ] cookie support - [ ] fancy argument parsing - [ ] document new classes (docstrings, maybe comments) - [ ] write examples @@ -74,3 +73,4 @@ i. e. the text between the end of the "!echo" and the end of the whole message. yaboli using `pip install git+https://github.com/Garmelon/yaboli`) - [x] implement !restart - [x] write project readme +- [x] cookie support diff --git a/yaboli/bot.py b/yaboli/bot.py index d08bcdd..3006773 100644 --- a/yaboli/bot.py +++ b/yaboli/bot.py @@ -33,9 +33,16 @@ class Bot(Client): nick = self.config[self.GENERAL_SECTION].get("nick") if nick is None: - logger.warn("No nick set in config file. Defaulting to empty nick") + logger.warn(("'nick' not set in config file. Defaulting to empty" + " nick")) nick = "" - super().__init__(nick) + + cookie_file = self.config[self.GENERAL_SECTION].get("cookie_file") + if cookie_file is None: + logger.warn(("'cookie_file' not set in config file. Using no cookie" + " file.")) + + super().__init__(nick, cookie_file=cookie_file) self._commands: List[Command] = [] diff --git a/yaboli/client.py b/yaboli/client.py index 5117c45..46777b0 100644 --- a/yaboli/client.py +++ b/yaboli/client.py @@ -1,7 +1,7 @@ import asyncio import functools import logging -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Union from .message import LiveMessage from .room import Room @@ -12,8 +12,12 @@ logger = logging.getLogger(__name__) __all__ = ["Client"] class Client: - def __init__(self, default_nick: str) -> None: + def __init__(self, + default_nick: str, + cookie_file: Optional[str] = None, + ) -> None: self._default_nick = default_nick + self._cookie_file = cookie_file self._rooms: Dict[str, List[Room]] = {} self._stop = asyncio.Event() @@ -48,13 +52,31 @@ class Client: async def join(self, room_name: str, password: Optional[str] = None, - nick: Optional[str] = None + nick: Optional[str] = None, + cookie_file: Union[str, bool] = True, ) -> Optional[Room]: + """ + cookie_file is the name of the file to store the cookies in. If it is + True, the client default is used. If it is False, no cookie file name + will be used. + """ + logger.info(f"Joining &{room_name}") if nick is None: nick = self._default_nick - room = Room(room_name, password=password, target_nick=nick) + + this_cookie_file: Optional[str] + + if isinstance(cookie_file, str): # This way, mypy doesn't complain + this_cookie_file = cookie_file + elif cookie_file: + this_cookie_file = self._cookie_file + else: + this_cookie_file = None + + room = Room(room_name, password=password, target_nick=nick, + cookie_file=this_cookie_file) room.register_event("connected", functools.partial(self.on_connected, room)) diff --git a/yaboli/connection.py b/yaboli/connection.py index fbb354f..af31d1c 100644 --- a/yaboli/connection.py +++ b/yaboli/connection.py @@ -6,6 +6,7 @@ from typing import Any, Awaitable, Callable, Dict, Optional import websockets +from .cookiejar import CookieJar from .events import Events from .exceptions import * @@ -97,8 +98,9 @@ class Connection: # Initialising - def __init__(self, url: str) -> None: + def __init__(self, url: str, cookie_file: Optional[str] = None) -> None: self._url = url + self._cookie_jar = CookieJar(cookie_file) self._events = Events() self._packet_id = 0 @@ -181,7 +183,8 @@ class Connection: try: logger.debug(f"Creating ws connection to {self._url!r}") - ws = await websockets.connect(self._url) + ws = await websockets.connect(self._url, + extra_headers=self._cookie_jar.get_cookies_as_headers()) self._ws = ws self._awaiting_replies = {} @@ -189,6 +192,11 @@ class Connection: self._ping_check = asyncio.create_task( self._disconnect_in(self.PING_TIMEOUT)) + # Put received cookies into cookie jar + for set_cookie in ws.response_headers.get_all("Set-Cookie"): + self._cookie_jar.add_cookie(set_cookie) + self._cookie_jar.save() + return True except (websockets.InvalidHandshake, websockets.InvalidStatusCode, diff --git a/yaboli/cookiejar.py b/yaboli/cookiejar.py new file mode 100644 index 0000000..833dbcb --- /dev/null +++ b/yaboli/cookiejar.py @@ -0,0 +1,77 @@ +import contextlib +import http.cookies as cookies +import logging +from typing import List, Optional, Tuple + +logger = logging.getLogger(__name__) + +__all__ = ["CookieJar"] + +class CookieJar: + """ + Keeps your cookies in a file. + + CookieJar doesn't attempt to discard old cookies, but that doesn't appear + to be necessary for keeping euphoria session cookies. + """ + + def __init__(self, filename: Optional[str] = None) -> None: + self._filename = filename + self._cookies = cookies.SimpleCookie() + + if not self._filename: + logger.warning("Could not load cookies, no filename given.") + return + + with contextlib.suppress(FileNotFoundError): + logger.info(f"Loading cookies from {self._filename!r}") + with open(self._filename, "r") as f: + for line in f: + self._cookies.load(line) + + def get_cookies(self) -> List[str]: + return [morsel.OutputString(attrs=[]) + for morsel in self._cookies.values()] + + def get_cookies_as_headers(self) -> List[Tuple[str, str]]: + """ + Return all stored cookies as tuples in a list. The first tuple entry is + always "Cookie". + """ + + return [("Cookie", cookie) for cookie in self.get_cookies()] + + def add_cookie(self, cookie: str) -> None: + """ + Parse cookie and add it to the jar. + + Example cookie: "a=bcd; Path=/; Expires=Wed, 24 Jul 2019 14:57:52 GMT; + HttpOnly; Secure" + """ + + logger.debug(f"Adding cookie {cookie!r}") + self._cookies.load(cookie) + + def save(self) -> None: + """ + Saves all current cookies to the cookie jar file. + """ + + if not self._filename: + logger.warning("Could not save cookies, no filename given.") + return + + logger.info(f"Saving cookies to {self._filename!r}") + + with open(self._filename, "w") as f: + for morsel in self._cookies.values(): + cookie_string = morsel.OutputString() + f.write(f"{cookie_string}\n") + + def clear(self) -> None: + """ + Removes all cookies from the cookie jar. + """ + + logger.debug("OMNOMNOM, cookies are all gone!") + self._cookies = cookies.SimpleCookie() diff --git a/yaboli/room.py b/yaboli/room.py index 4196d5f..053e398 100644 --- a/yaboli/room.py +++ b/yaboli/room.py @@ -60,7 +60,8 @@ class Room: name: str, password: Optional[str] = None, target_nick: str = "", - url_format: str = URL_FORMAT + url_format: str = URL_FORMAT, + cookie_file: Optional[str] = None, ) -> None: self._name = name self._password = password @@ -78,7 +79,7 @@ class Room: # Connected management self._url = self._url_format.format(self._name) - self._connection = Connection(self._url) + self._connection = Connection(self._url, cookie_file=cookie_file) self._events = Events() self._connected = asyncio.Event()