| OLD | NEW |
| (Empty) |
| 1 # Copyright 2014 Google Inc. All rights reserved. | |
| 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, overwrite=False): | |
| 120 """Write a Credentials to the datastore. | |
| 121 | |
| 122 Args: | |
| 123 credentials: Credentials, the credentials to store. | |
| 124 overwrite: Boolean, indicates whether you would like these credentials to | |
| 125 overwrite any existing stored credentials. | |
| 126 """ | |
| 127 args = {self.key_name: self.key_value} | |
| 128 | |
| 129 if overwrite: | |
| 130 entity, unused_is_new = self.model_class.objects.get_or_create(**args) | |
| 131 else: | |
| 132 entity = self.model_class(**args) | |
| 133 | |
| 134 setattr(entity, self.property_name, credentials) | |
| 135 entity.save() | |
| 136 | |
| 137 def locked_delete(self): | |
| 138 """Delete Credentials from the datastore.""" | |
| 139 | |
| 140 query = {self.key_name: self.key_value} | |
| 141 entities = self.model_class.objects.filter(**query).delete() | |
| OLD | NEW |