| Adam Israel | c3e6c2e | 2018-03-01 09:31:50 -0500 | [diff] [blame] | 1 | import datetime |
| 2 | import http.cookiejar as cookiejar |
| 3 | import json |
| 4 | import time |
| 5 | |
| 6 | import pyrfc3339 |
| 7 | |
| 8 | |
| 9 | class GoCookieJar(cookiejar.FileCookieJar): |
| 10 | '''A CookieJar implementation that reads and writes cookies |
| 11 | to the cookiejar format as understood by the Go package |
| 12 | github.com/juju/persistent-cookiejar.''' |
| 13 | def _really_load(self, f, filename, ignore_discard, ignore_expires): |
| 14 | '''Implement the _really_load method called by FileCookieJar |
| 15 | to implement the actual cookie loading''' |
| 16 | data = json.load(f) or [] |
| 17 | now = time.time() |
| 18 | for cookie in map(_new_py_cookie, data): |
| 19 | if not ignore_expires and cookie.is_expired(now): |
| 20 | continue |
| 21 | self.set_cookie(cookie) |
| 22 | |
| 23 | def save(self, filename=None, ignore_discard=False, ignore_expires=False): |
| 24 | '''Implement the FileCookieJar abstract method.''' |
| 25 | if filename is None: |
| 26 | if self.filename is not None: |
| 27 | filename = self.filename |
| 28 | else: |
| 29 | raise ValueError(cookiejar.MISSING_FILENAME_TEXT) |
| 30 | |
| 31 | # TODO: obtain file lock, read contents of file, and merge with |
| 32 | # current content. |
| 33 | go_cookies = [] |
| 34 | now = time.time() |
| 35 | for cookie in self: |
| 36 | if not ignore_discard and cookie.discard: |
| 37 | continue |
| 38 | if not ignore_expires and cookie.is_expired(now): |
| 39 | continue |
| 40 | go_cookies.append(_new_go_cookie(cookie)) |
| 41 | with open(filename, "w") as f: |
| 42 | f.write(json.dumps(go_cookies)) |
| 43 | |
| 44 | |
| 45 | def _new_py_cookie(go_cookie): |
| 46 | '''Convert a Go-style JSON-unmarshaled cookie into a Python cookie''' |
| 47 | expires = None |
| 48 | if go_cookie.get('Expires') is not None: |
| 49 | t = pyrfc3339.parse(go_cookie['Expires']) |
| 50 | expires = t.timestamp() |
| 51 | return cookiejar.Cookie( |
| 52 | version=0, |
| 53 | name=go_cookie['Name'], |
| 54 | value=go_cookie['Value'], |
| 55 | port=None, |
| 56 | port_specified=False, |
| 57 | # Unfortunately Python cookies don't record the original |
| 58 | # host that the cookie came from, so we'll just use Domain |
| 59 | # for that purpose, and record that the domain was specified, |
| 60 | # even though it probably was not. This means that |
| 61 | # we won't correctly record the CanonicalHost entry |
| 62 | # when writing the cookie file after reading it. |
| 63 | domain=go_cookie['Domain'], |
| 64 | domain_specified=not go_cookie['HostOnly'], |
| 65 | domain_initial_dot=False, |
| 66 | path=go_cookie['Path'], |
| 67 | path_specified=True, |
| 68 | secure=go_cookie['Secure'], |
| 69 | expires=expires, |
| 70 | discard=False, |
| 71 | comment=None, |
| 72 | comment_url=None, |
| 73 | rest=None, |
| 74 | rfc2109=False, |
| 75 | ) |
| 76 | |
| 77 | |
| 78 | def _new_go_cookie(py_cookie): |
| 79 | '''Convert a python cookie to the JSON-marshalable Go-style cookie form.''' |
| 80 | # TODO (perhaps): |
| 81 | # HttpOnly |
| 82 | # Creation |
| 83 | # LastAccess |
| 84 | # Updated |
| 85 | # not done properly: CanonicalHost. |
| 86 | go_cookie = { |
| 87 | 'Name': py_cookie.name, |
| 88 | 'Value': py_cookie.value, |
| 89 | 'Domain': py_cookie.domain, |
| 90 | 'HostOnly': not py_cookie.domain_specified, |
| 91 | 'Persistent': not py_cookie.discard, |
| 92 | 'Secure': py_cookie.secure, |
| 93 | 'CanonicalHost': py_cookie.domain, |
| 94 | } |
| 95 | if py_cookie.path_specified: |
| 96 | go_cookie['Path'] = py_cookie.path |
| 97 | if py_cookie.expires is not None: |
| 98 | unix_time = datetime.datetime.fromtimestamp(py_cookie.expires) |
| 99 | # Note: fromtimestamp bizarrely produces a time without |
| 100 | # a time zone, so we need to use accept_naive. |
| 101 | go_cookie['Expires'] = pyrfc3339.generate(unix_time, accept_naive=True) |
| 102 | return go_cookie |