-
Notifications
You must be signed in to change notification settings - Fork 2
/
tom_client.py
94 lines (72 loc) · 4.08 KB
/
tom_client.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
import requests
class TomClient:
"""A thin class that supports sending requests via "requests" to the DESC tom.
Usage: initialize one of these, giving it the url, your TOM
username, and either your TOM password, or a file that has your TOM
password in it:
tc = TomClient( username='rknop', passwordfile='/home/raknop/secrets/tom_rknop_passwd' )
(You can give it a url with url=; it defaults to https://desc-tom.lbl.gov.)
Thereafter, just do something like
res = tc.request( "POST", "elasticc2/ppdbdiaobject/55772173" )
and res will come back with a string that you can load into JSON
that will have all the fun data about PPDBDiaObject number 55772173.
tc.request is just a thin front-end to python requests.request. The
only reason to use this client rather than the python requests
module directly is that this class takes care of the stupid fiddly
bits of getting some headers that django demands set up right in the
request object when you log in.
"""
def __init__( self, url="https://desc-tom.lbl.gov", username=None, password=None, passwordfile=None, connect=True ):
self._url = url
self._username = username
self._password = password
self._rqs = None
if self._password is None:
if passwordfile is None:
raise RuntimeError( "Must give either password or passwordfile. " )
with open( passwordfile ) as ifp:
self._password = ifp.readline().strip()
if connect:
self.connect()
def connect( self ):
self._rqs = requests.session()
res = self._rqs.get( f'{self._url}/accounts/login/' )
if res.status_code != 200:
raise RuntimeError( f"Got status {res.status_code} from first attempt to connect to {self._url}" )
res = self._rqs.post( f'{self._url}/accounts/login/',
data={ 'username': self._username,
'password': self._password,
'csrfmiddlewaretoken': self._rqs.cookies['csrftoken'] } )
if res.status_code != 200:
raise RuntimeError( f"Failed to log in; http status: {res.status_code}" )
if 'Please enter a correct' in res.text:
# This is a very cheesy attempt at checking if the login failed.
# I haven't found clean documentation on how to log into a django site
# from an app like this using standard authentication stuff. So, for
# now, I'm counting on the HTML that happened to come back when
# I ran it with a failed login one time. One of these days I'll actually
# figure out how Django auth works and make a version of /accounts/login/
# designed for use in API scripts like this one, rather than desgined
# for interactive users.
raise RuntimeError( "Failed to log in. I think. Put in a debug break and look at res.text" )
self._rqs.headers.update( { 'X-CSRFToken': self._rqs.cookies['csrftoken'] } )
def request( self, method="GET", page=None, **kwargs ):
"""Send a request to the TOM
method : a string with the HTTP method ("GET", "POST", etc.)
page : the page to get; this is the URL but with the url you
passed to the constructor removed. So, if you wanted to get
https://desc-tom.lbl.gov/elasticc, you'd pass just "elasticc"
here.
**kwargs : additional keyword arguments are passed on to
requests.request
"""
return self._rqs.request( method=method, url=f"{self._url}/{page}", **kwargs )
def post( self, page=None, **kwargs ):
"""Shortand for TomClient.request( "POST", ... )"""
return self.request( "POST", page, **kwargs )
def get( self, page=None, **kwargs ):
"""Shortand for TomClient.request( "GET", ... )"""
return self.request( "GET", page, **kwargs )
def put( self, page=None, **kwargs ):
"""Shortand for TomClient.request( "PUT", ... )"""
return self.request( "PUT", page, **kwargs )