storage.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224
  1. """Django ORM models for Social Auth"""
  2. import base64
  3. import six
  4. import sys
  5. from django.core.exceptions import FieldDoesNotExist
  6. from django.db import transaction, router
  7. from django.db.utils import IntegrityError
  8. from social_core.storage import UserMixin, AssociationMixin, NonceMixin, \
  9. CodeMixin, PartialMixin, BaseStorage
  10. class DjangoUserMixin(UserMixin):
  11. """Social Auth association model"""
  12. @classmethod
  13. def changed(cls, user):
  14. user.save()
  15. def set_extra_data(self, extra_data=None):
  16. if super(DjangoUserMixin, self).set_extra_data(extra_data):
  17. self.save()
  18. @classmethod
  19. def allowed_to_disconnect(cls, user, backend_name, association_id=None):
  20. if association_id is not None:
  21. qs = cls.objects.exclude(id=association_id)
  22. else:
  23. qs = cls.objects.exclude(provider=backend_name)
  24. qs = qs.filter(user=user)
  25. if hasattr(user, 'has_usable_password'):
  26. valid_password = user.has_usable_password()
  27. else:
  28. valid_password = True
  29. return valid_password or qs.count() > 0
  30. @classmethod
  31. def disconnect(cls, entry):
  32. entry.delete()
  33. @classmethod
  34. def username_field(cls):
  35. return getattr(cls.user_model(), 'USERNAME_FIELD', 'username')
  36. @classmethod
  37. def user_exists(cls, *args, **kwargs):
  38. """
  39. Return True/False if a User instance exists with the given arguments.
  40. Arguments are directly passed to filter() manager method.
  41. """
  42. if 'username' in kwargs:
  43. kwargs[cls.username_field()] = kwargs.pop('username')
  44. return cls.user_model().objects.filter(*args, **kwargs).count() > 0
  45. @classmethod
  46. def get_username(cls, user):
  47. return getattr(user, cls.username_field(), None)
  48. @classmethod
  49. def create_user(cls, *args, **kwargs):
  50. username_field = cls.username_field()
  51. if 'username' in kwargs:
  52. if username_field not in kwargs:
  53. kwargs[username_field] = kwargs.pop('username')
  54. else:
  55. # If username_field is 'email' and there is no field named "username"
  56. # then latest should be removed from kwargs.
  57. try:
  58. cls.user_model()._meta.get_field('username')
  59. except FieldDoesNotExist:
  60. kwargs.pop('username')
  61. try:
  62. if hasattr(transaction, 'atomic'):
  63. # In Django versions that have an "atomic" transaction decorator / context
  64. # manager, there's a transaction wrapped around this call.
  65. # If the create fails below due to an IntegrityError, ensure that the transaction
  66. # stays undamaged by wrapping the create in an atomic.
  67. using = router.db_for_write(cls.user_model())
  68. with transaction.atomic(using=using):
  69. user = cls.user_model().objects.create_user(*args, **kwargs)
  70. else:
  71. user = cls.user_model().objects.create_user(*args, **kwargs)
  72. except IntegrityError:
  73. # User might have been created on a different thread, try and find them.
  74. # If we don't, re-raise the IntegrityError.
  75. exc_info = sys.exc_info()
  76. # If email comes in as None it won't get found in the get
  77. if kwargs.get('email', True) is None:
  78. kwargs['email'] = ''
  79. try:
  80. user = cls.user_model().objects.get(*args, **kwargs)
  81. except cls.user_model().DoesNotExist:
  82. six.reraise(*exc_info)
  83. return user
  84. @classmethod
  85. def get_user(cls, pk=None, **kwargs):
  86. if pk:
  87. kwargs = {'pk': pk}
  88. try:
  89. return cls.user_model().objects.get(**kwargs)
  90. except cls.user_model().DoesNotExist:
  91. return None
  92. @classmethod
  93. def get_users_by_email(cls, email):
  94. user_model = cls.user_model()
  95. email_field = getattr(user_model, 'EMAIL_FIELD', 'email')
  96. return user_model.objects.filter(**{email_field + '__iexact': email})
  97. @classmethod
  98. def get_social_auth(cls, provider, uid):
  99. if not isinstance(uid, six.string_types):
  100. uid = str(uid)
  101. try:
  102. return cls.objects.get(provider=provider, uid=uid)
  103. except cls.DoesNotExist:
  104. return None
  105. @classmethod
  106. def get_social_auth_for_user(cls, user, provider=None, id=None):
  107. qs = cls.objects.filter(user=user)
  108. if provider:
  109. qs = qs.filter(provider=provider)
  110. if id:
  111. qs = qs.filter(id=id)
  112. return qs
  113. @classmethod
  114. def create_social_auth(cls, user, uid, provider):
  115. if not isinstance(uid, six.string_types):
  116. uid = str(uid)
  117. if hasattr(transaction, 'atomic'):
  118. # In Django versions that have an "atomic" transaction decorator / context
  119. # manager, there's a transaction wrapped around this call.
  120. # If the create fails below due to an IntegrityError, ensure that the transaction
  121. # stays undamaged by wrapping the create in an atomic.
  122. using = router.db_for_write(cls)
  123. with transaction.atomic(using=using):
  124. social_auth = cls.objects.create(user=user, uid=uid, provider=provider)
  125. else:
  126. social_auth = cls.objects.create(user=user, uid=uid, provider=provider)
  127. return social_auth
  128. class DjangoNonceMixin(NonceMixin):
  129. @classmethod
  130. def use(cls, server_url, timestamp, salt):
  131. return cls.objects.get_or_create(server_url=server_url,
  132. timestamp=timestamp,
  133. salt=salt)[1]
  134. @classmethod
  135. def get(cls, server_url, salt):
  136. return cls.objects.get(
  137. server_url=server_url,
  138. salt=salt,
  139. )
  140. @classmethod
  141. def delete(cls, nonce):
  142. nonce.delete()
  143. class DjangoAssociationMixin(AssociationMixin):
  144. @classmethod
  145. def store(cls, server_url, association):
  146. # Don't use get_or_create because issued cannot be null
  147. try:
  148. assoc = cls.objects.get(server_url=server_url,
  149. handle=association.handle)
  150. except cls.DoesNotExist:
  151. assoc = cls(server_url=server_url,
  152. handle=association.handle)
  153. try:
  154. assoc.secret = base64.encodebytes(association.secret).decode()
  155. except AttributeError:
  156. assoc.secret = base64.encodestring(association.secret).decode()
  157. assoc.issued = association.issued
  158. assoc.lifetime = association.lifetime
  159. assoc.assoc_type = association.assoc_type
  160. assoc.save()
  161. @classmethod
  162. def get(cls, *args, **kwargs):
  163. return cls.objects.filter(*args, **kwargs)
  164. @classmethod
  165. def remove(cls, ids_to_delete):
  166. cls.objects.filter(pk__in=ids_to_delete).delete()
  167. class DjangoCodeMixin(CodeMixin):
  168. @classmethod
  169. def get_code(cls, code):
  170. try:
  171. return cls.objects.get(code=code)
  172. except cls.DoesNotExist:
  173. return None
  174. class DjangoPartialMixin(PartialMixin):
  175. @classmethod
  176. def load(cls, token):
  177. try:
  178. return cls.objects.get(token=token)
  179. except cls.DoesNotExist:
  180. return None
  181. @classmethod
  182. def destroy(cls, token):
  183. partial = cls.load(token)
  184. if partial:
  185. partial.delete()
  186. class BaseDjangoStorage(BaseStorage):
  187. user = DjangoUserMixin
  188. nonce = DjangoNonceMixin
  189. association = DjangoAssociationMixin
  190. code = DjangoCodeMixin