From c6260a491fc4d54f129033014536bd7556ec1c65 Mon Sep 17 00:00:00 2001
From: Chase Geigle <geigle1@illinois.edu>
Date: Mon, 6 Aug 2018 18:32:16 -0500
Subject: [PATCH] fix(dev): Use requests.Session object for updating cookies

Before, we were not keeping track of cookie updates being sent by the
Piazza API when making requests. This is normally not an issue, but
presents a problem when a machine is making multiple requests to the
Piazza API for different networks. Imagine running a crawler for two
courses: when the second crawler begins crawling a different course, the
sessions for both crawlers become invalid because the server must be
sending updated cookies for them both in response to content.get calls.
By using a session object to make our requests, we can more simply
ensure that the cookies stay up to date, and this fixes the issues with
multiple crawlers on one machine.
---
 piazza_api/network.py |  7 ++++---
 piazza_api/piazza.py  |  2 +-
 piazza_api/rpc.py     | 21 +++++++++------------
 3 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/piazza_api/network.py b/piazza_api/network.py
index 9d68705..d4c9b6d 100644
--- a/piazza_api/network.py
+++ b/piazza_api/network.py
@@ -47,12 +47,13 @@ class Network(object):
     """Abstraction for a Piazza "Network" (or class)
 
     :param network_id: ID of the network
-    :param cookies: RequestsCookieJar containing cookies used for authentication
+    :param session: requests.Session object containing cookies used for
+        authentication
     """
-    def __init__(self, network_id, cookies):
+    def __init__(self, network_id, session):
         self._nid = network_id
         self._rpc = PiazzaRPC(network_id=self._nid)
-        self._rpc.cookies = cookies
+        self._rpc.session = session
 
         ff = namedtuple('FeedFilters', ['unread', 'following', 'folder'])
         self._feed_filters = ff(UnreadFilter, FollowingFilter, FolderFilter)
diff --git a/piazza_api/piazza.py b/piazza_api/piazza.py
index 98a89f5..ee6c8df 100644
--- a/piazza_api/piazza.py
+++ b/piazza_api/piazza.py
@@ -43,7 +43,7 @@ def network(self, network_id):
             https://piazza.com/class/{network_id}
         """
         self._ensure_authenticated()
-        return Network(network_id, self._rpc_api.cookies)
+        return Network(network_id, self._rpc_api.session)
 
     def get_user_profile(self):
         """Get profile of the current user
diff --git a/piazza_api/rpc.py b/piazza_api/rpc.py
index 133f0df..a7e4a73 100644
--- a/piazza_api/rpc.py
+++ b/piazza_api/rpc.py
@@ -31,7 +31,7 @@ def __init__(self, network_id=None):
             "logic": "https://piazza.com/logic/api",
             "main": "https://piazza.com/main/api",
         }
-        self.cookies = None
+        self.session = requests.Session()
 
     def user_login(self, email=None, password=None):
         """Login with email, password and get back a session cookie
@@ -51,14 +51,13 @@ def user_login(self, email=None, password=None):
         }
         # If the user/password match, the server respond will contain a
         #  session cookie that you can use to authenticate future requests.
-        r = requests.post(
+        r = self.session.post(
             self.base_api_urls["logic"],
             data=json.dumps(login_data),
         )
         if r.json()["result"] not in ["OK"]:
             raise AuthenticationError("Could not authenticate.\n{}"
                                       .format(r.json()))
-        self.cookies = r.cookies
 
     def demo_login(self, auth=None, url=None):
         """Authenticate with a "Share Your Class" URL using a demo user.
@@ -76,10 +75,9 @@ def demo_login(self, auth=None, url=None):
         if url is None:
             url = "https://piazza.com/demo_login"
             params = dict(nid=self._nid, auth=auth)
-            res = requests.get(url, params=params)
+            res = self.session.get(url, params=params)
         else:
-            res = requests.get(url)
-        self.cookies = res.cookies
+            res = self.session.get(url)
 
     def content_get(self, cid, nid=None):
         """Get data from post `cid` in network `nid`
@@ -379,9 +377,9 @@ def request(self, method, data=None, nid=None, nid_key='nid',
             data = {}
 
         headers = {}
-        if "session_id" in self.cookies:
-            headers["CSRF-Token"] = self.cookies["session_id"]
-        
+        if "session_id" in self.session.cookies:
+            headers["CSRF-Token"] = self.session.cookies["session_id"]
+
         # Adding a nonce to the request
         endpoint = self.base_api_urls[api_type]
         if api_type == "logic":
@@ -390,13 +388,12 @@ def request(self, method, data=None, nid=None, nid_key='nid',
                 _piazza_nonce()
             )
 
-        response = requests.post(
+        response = self.session.post(
             endpoint,
             data=json.dumps({
                 "method": method,
                 "params": dict({nid_key: nid}, **data)
             }),
-            cookies=self.cookies,
             headers=headers
         )
         return response if return_response else response.json()
@@ -410,7 +407,7 @@ def _check_authenticated(self):
 
         :raises: NotAuthenticatedError
         """
-        if self.cookies is None:
+        if not self.session.cookies:
             raise NotAuthenticatedError("You must authenticate before "
                                         "making any other requests.")