1# Copyright 2016 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"""Contains classes used for the Django ORM storage."""
16
17import base64
18import pickle
19
20from django.db import models
21from django.utils import encoding
22
23import oauth2client
24
25
26class CredentialsField(models.Field):
27    """Django ORM field for storing OAuth2 Credentials."""
28
29    def __init__(self, *args, **kwargs):
30        if 'null' not in kwargs:
31            kwargs['null'] = True
32        super(CredentialsField, self).__init__(*args, **kwargs)
33
34    def get_internal_type(self):
35        return 'BinaryField'
36
37    def from_db_value(self, value, expression, connection, context):
38        """Overrides ``models.Field`` method. This converts the value
39        returned from the database to an instance of this class.
40        """
41        return self.to_python(value)
42
43    def to_python(self, value):
44        """Overrides ``models.Field`` method. This is used to convert
45        bytes (from serialization etc) to an instance of this class"""
46        if value is None:
47            return None
48        elif isinstance(value, oauth2client.client.Credentials):
49            return value
50        else:
51            return pickle.loads(base64.b64decode(encoding.smart_bytes(value)))
52
53    def get_prep_value(self, value):
54        """Overrides ``models.Field`` method. This is used to convert
55        the value from an instances of this class to bytes that can be
56        inserted into the database.
57        """
58        if value is None:
59            return None
60        else:
61            return encoding.smart_text(base64.b64encode(pickle.dumps(value)))
62
63    def value_to_string(self, obj):
64        """Convert the field value from the provided model to a string.
65
66        Used during model serialization.
67
68        Args:
69            obj: db.Model, model object
70
71        Returns:
72            string, the serialized field value
73        """
74        value = self._get_val_from_obj(obj)
75        return self.get_prep_value(value)
76