Add cookie support

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Next version
 
+- add cookie support
 - add !restart to botrulez
 - save (overwrite) `Bot` config file
 

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

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] = []
 

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))

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,

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()

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()