importer.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198
  1. from typing import TYPE_CHECKING
  2. import aiohttp
  3. import tablib
  4. from asgiref.sync import sync_to_async
  5. from django.db.models import Q
  6. from django.urls import reverse
  7. from apps.organizations_ext.models import OrganizationUser, OrganizationUserRole
  8. from apps.organizations_ext.resources import (
  9. OrganizationResource,
  10. OrganizationUserResource,
  11. )
  12. from apps.projects.models import Project
  13. from apps.projects.resources import ProjectKeyResource, ProjectResource
  14. from apps.teams.resources import TeamResource
  15. from apps.users.models import User
  16. from apps.users.resources import UserResource
  17. from .exceptions import ImporterException
  18. if TYPE_CHECKING:
  19. from apps.shared.types import TypeJson
  20. class GlitchTipImporter:
  21. """
  22. Generic importer tool to use with cli or web
  23. If used by a non server admin, it's important to assume all incoming
  24. JSON is hostile and not from a real GT server. Foreign Key ids could be
  25. faked and used to elevate privileges. Always confirm new data is associated with
  26. appropriate organization. Also assume user is at least an org admin, no need to
  27. double check permissions when creating assets within the organization.
  28. create_users should be False unless running as superuser/management command
  29. """
  30. def __init__(
  31. self, url: str, auth_token: str, organization_slug: str, create_users=False
  32. ):
  33. self.url = url.rstrip("/")
  34. self.headers = {"Authorization": f"Bearer {auth_token}"}
  35. self.create_users = create_users
  36. self.organization_slug = organization_slug
  37. self.organization_id = None
  38. self.organization_url = reverse(
  39. "api:get_organization", args=[self.organization_slug]
  40. )
  41. self.organization_users_url = reverse(
  42. "api:list_organization_members",
  43. kwargs={"organization_slug": self.organization_slug},
  44. )
  45. self.projects_url = reverse(
  46. "api:list_organization_projects", args=[self.organization_slug]
  47. )
  48. self.teams_url = reverse("api:list_teams", args=[self.organization_slug])
  49. async def run(self, organization_id=None):
  50. """Set organization_id to None to import (superuser only)"""
  51. if organization_id is None:
  52. await self.import_organization()
  53. else:
  54. self.organization_id = organization_id
  55. await self.import_organization_users()
  56. await self.import_projects()
  57. await self.import_teams()
  58. async def get(self, url: str) -> "TypeJson":
  59. async with aiohttp.ClientSession() as session:
  60. async with session.get(url, headers=self.headers) as res:
  61. return await res.json()
  62. async def import_organization(self):
  63. resource = OrganizationResource()
  64. data = await self.get(self.url + self.organization_url)
  65. self.organization_id = data["id"] # TODO unsafe for web usage
  66. dataset = tablib.Dataset()
  67. dataset.dict = [data]
  68. await sync_to_async(resource.import_data)(dataset, raise_errors=True)
  69. async def import_organization_users(self):
  70. resource = OrganizationUserResource()
  71. org_users = await self.get(self.url + self.organization_users_url)
  72. if not org_users:
  73. return
  74. if self.create_users:
  75. user_resource = UserResource()
  76. users_list = [
  77. org_user["user"] for org_user in org_users if org_user is not None
  78. ]
  79. users = [
  80. {k: v for k, v in user.items() if k in ["id", "email", "name"]}
  81. for user in users_list
  82. ]
  83. dataset = tablib.Dataset()
  84. dataset.dict = users
  85. await sync_to_async(user_resource.import_data)(dataset, raise_errors=True)
  86. for org_user in org_users:
  87. org_user["organization"] = self.organization_id
  88. org_user["role"] = OrganizationUserRole.from_string(org_user["role"])
  89. if self.create_users:
  90. org_user["user"] = (
  91. User.objects.filter(email=org_user["user"]["email"])
  92. .values_list("pk", flat=True)
  93. .first()
  94. )
  95. else:
  96. org_user["user"] = None
  97. dataset = tablib.Dataset()
  98. dataset.dict = org_users
  99. await sync_to_async(resource.import_data)(dataset, raise_errors=True)
  100. async def import_projects(self):
  101. project_resource = ProjectResource()
  102. project_key_resource = ProjectKeyResource()
  103. projects = await self.get(self.url + self.projects_url)
  104. project_keys = []
  105. for project in projects:
  106. project["organization"] = self.organization_id
  107. keys = await self.get(
  108. self.url
  109. + reverse(
  110. "api:list_project_keys",
  111. args=[self.organization_slug, project["slug"]],
  112. )
  113. )
  114. for key in keys:
  115. key["project"] = project["id"]
  116. key["public_key"] = key["public"]
  117. project_keys += keys
  118. dataset = tablib.Dataset()
  119. dataset.dict = projects
  120. await sync_to_async(project_resource.import_data)(dataset, raise_errors=True)
  121. owned_project_ids = [
  122. pk
  123. async for pk in Project.objects.filter(
  124. organization_id=self.organization_id,
  125. pk__in=[d["projectId"] for d in project_keys],
  126. ).values_list("pk", flat=True)
  127. ]
  128. project_keys = list(
  129. filter(lambda key: key["projectId"] in owned_project_ids, project_keys)
  130. )
  131. dataset.dict = project_keys
  132. await sync_to_async(project_key_resource.import_data)(
  133. dataset, raise_errors=True
  134. )
  135. async def import_teams(self):
  136. resource = TeamResource()
  137. teams = await self.get(self.url + self.teams_url)
  138. for team in teams:
  139. team["organization"] = self.organization_id
  140. team["projects"] = ",".join(
  141. map(
  142. str,
  143. [
  144. pk
  145. async for pk in Project.objects.filter(
  146. organization_id=self.organization_id,
  147. pk__in=[int(d["id"]) for d in team["projects"]],
  148. ).values_list("id", flat=True)
  149. ],
  150. )
  151. )
  152. team_members = await self.get(
  153. self.url
  154. + reverse(
  155. "api:list_team_organization_members",
  156. args=[self.organization_slug, team["slug"]],
  157. )
  158. )
  159. team_member_emails = [d["email"] for d in team_members]
  160. team["members"] = ",".join(
  161. [
  162. str(i)
  163. async for i in OrganizationUser.objects.filter(
  164. organization_id=self.organization_id
  165. )
  166. .filter(
  167. Q(email__in=team_member_emails)
  168. | Q(user__email__in=team_member_emails)
  169. )
  170. .values_list("pk", flat=True)
  171. ]
  172. )
  173. dataset = tablib.Dataset()
  174. dataset.dict = teams
  175. await sync_to_async(resource.import_data)(dataset, raise_errors=True)
  176. async def check_auth(self):
  177. async with aiohttp.ClientSession() as session:
  178. async with session.get(self.url + "/api/0/", headers=self.headers) as res:
  179. data = await res.json()
  180. if res.status != 200 or not data["user"]:
  181. raise ImporterException("Bad auth token")