OLD | NEW |
| (Empty) |
1 # Copyright (C) 2010 Google Inc. | |
2 # | |
3 # Licensed under the Apache License, Version 2.0 (the "License"); | |
4 # you may not use this file except in compliance with the License. | |
5 # You may obtain a copy of the License at | |
6 # | |
7 # http://www.apache.org/licenses/LICENSE-2.0 | |
8 # | |
9 # Unless required by applicable law or agreed to in writing, software | |
10 # distributed under the License is distributed on an "AS IS" BASIS, | |
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
12 # See the License for the specific language governing permissions and | |
13 # limitations under the License. | |
14 | |
15 """OAuth 2.0 utilities for Django. | |
16 | |
17 Utilities for using OAuth 2.0 in conjunction with | |
18 the Django datastore. | |
19 """ | |
20 | |
21 __author__ = 'jcgregorio@google.com (Joe Gregorio)' | |
22 | |
23 import oauth2client | |
24 import base64 | |
25 import pickle | |
26 | |
27 from django.db import models | |
28 from oauth2client.client import Storage as BaseStorage | |
29 | |
30 class CredentialsField(models.Field): | |
31 | |
32 __metaclass__ = models.SubfieldBase | |
33 | |
34 def __init__(self, *args, **kwargs): | |
35 if 'null' not in kwargs: | |
36 kwargs['null'] = True | |
37 super(CredentialsField, self).__init__(*args, **kwargs) | |
38 | |
39 def get_internal_type(self): | |
40 return "TextField" | |
41 | |
42 def to_python(self, value): | |
43 if value is None: | |
44 return None | |
45 if isinstance(value, oauth2client.client.Credentials): | |
46 return value | |
47 return pickle.loads(base64.b64decode(value)) | |
48 | |
49 def get_db_prep_value(self, value, connection, prepared=False): | |
50 if value is None: | |
51 return None | |
52 return base64.b64encode(pickle.dumps(value)) | |
53 | |
54 | |
55 class FlowField(models.Field): | |
56 | |
57 __metaclass__ = models.SubfieldBase | |
58 | |
59 def __init__(self, *args, **kwargs): | |
60 if 'null' not in kwargs: | |
61 kwargs['null'] = True | |
62 super(FlowField, self).__init__(*args, **kwargs) | |
63 | |
64 def get_internal_type(self): | |
65 return "TextField" | |
66 | |
67 def to_python(self, value): | |
68 if value is None: | |
69 return None | |
70 if isinstance(value, oauth2client.client.Flow): | |
71 return value | |
72 return pickle.loads(base64.b64decode(value)) | |
73 | |
74 def get_db_prep_value(self, value, connection, prepared=False): | |
75 if value is None: | |
76 return None | |
77 return base64.b64encode(pickle.dumps(value)) | |
78 | |
79 | |
80 class Storage(BaseStorage): | |
81 """Store and retrieve a single credential to and from | |
82 the datastore. | |
83 | |
84 This Storage helper presumes the Credentials | |
85 have been stored as a CredenialsField | |
86 on a db model class. | |
87 """ | |
88 | |
89 def __init__(self, model_class, key_name, key_value, property_name): | |
90 """Constructor for Storage. | |
91 | |
92 Args: | |
93 model: db.Model, model class | |
94 key_name: string, key name for the entity that has the credentials | |
95 key_value: string, key value for the entity that has the credentials | |
96 property_name: string, name of the property that is an CredentialsProperty | |
97 """ | |
98 self.model_class = model_class | |
99 self.key_name = key_name | |
100 self.key_value = key_value | |
101 self.property_name = property_name | |
102 | |
103 def locked_get(self): | |
104 """Retrieve Credential from datastore. | |
105 | |
106 Returns: | |
107 oauth2client.Credentials | |
108 """ | |
109 credential = None | |
110 | |
111 query = {self.key_name: self.key_value} | |
112 entities = self.model_class.objects.filter(**query) | |
113 if len(entities) > 0: | |
114 credential = getattr(entities[0], self.property_name) | |
115 if credential and hasattr(credential, 'set_store'): | |
116 credential.set_store(self) | |
117 return credential | |
118 | |
119 def locked_put(self, credentials): | |
120 """Write a Credentials to the datastore. | |
121 | |
122 Args: | |
123 credentials: Credentials, the credentials to store. | |
124 """ | |
125 args = {self.key_name: self.key_value} | |
126 entity = self.model_class(**args) | |
127 setattr(entity, self.property_name, credentials) | |
128 entity.save() | |
129 | |
130 def locked_delete(self): | |
131 """Delete Credentials from the datastore.""" | |
132 | |
133 query = {self.key_name: self.key_value} | |
134 entities = self.model_class.objects.filter(**query).delete() | |
OLD | NEW |