b037d6bea61176829c5f1fdded38ab31a16bc3da
[bigly] / src / bigly / views.py
1 from urllib.parse import urlparse, urlunparse, parse_qs
2
3 from django.shortcuts import redirect, render
4 from django.views.generic.base import TemplateView
5
6 from rest_framework import status, viewsets
7 from rest_framework.response import Response
8 import requests
9
10 from . import serializers
11
12 def _remove_utm(link):
13     parsed_link = urlparse(link)
14     parsed_link = parsed_link._replace(
15         query='&'.join(
16             [
17                 p
18                 for p in parsed_link.query.split('&')
19                 if not p.startswith('utm_')
20             ]
21         )
22     )
23
24     return ''.join((
25         parsed_link.scheme + '://' if parsed_link.scheme else '',
26         parsed_link.netloc,
27         parsed_link.path,
28         ';' + parsed_link.params if parsed_link.params else '',
29         '?' + parsed_link.query if parsed_link.query else '',
30         '#' + parsed_link.fragment if parsed_link.fragment else '',
31     ))
32
33 def _follow_redirects(link, remove_utm):
34     redirect_sequence = []
35
36     while True:
37         redirect_sequence.append(link)
38
39         if remove_utm:
40             link = _remove_utm(link)
41
42         # TODO Do this in an async call so it doesn't block the main thread
43         response = requests.head(link, timeout=10)
44
45         # TODO Handle timeouts
46
47         if 301 <= response.status_code and response.status_code <= 308:
48             # TODO Handle the different kinds of redirects correctly
49
50             link = response.headers.get('Location')
51
52             if not link:
53                 # TODO Handle this
54                 raise Exception()
55
56         # TODO Handle error responses
57         else:
58             return {
59                 'link': link,
60                 'status': response.status_code,
61
62                 # TODO Handle different capitalizations of "Content-Type"
63                 'content_type': response.headers.get('Content-Type'),
64                 'redirect_sequence': redirect_sequence,
65             }
66
67 class IndexView(TemplateView):
68     template_name = 'bigly/index.html'
69
70 index = IndexView.as_view()
71
72 class FAQView(TemplateView):
73     template_name = 'bigly/faq.html'
74
75 faq = FAQView.as_view()
76
77 def embiggen(request):
78     serializer = serializers.FollowRedirectsSerializer(data=request.GET)
79
80     if not serializer.is_valid():
81         return render(
82             request,
83             'bigly/index.html',
84             {
85                 'errors': serializer.errors,
86             },
87             status=400,
88         )
89
90     result = _follow_redirects(
91         link = serializer.data['link'],
92         remove_utm = serializer.data['remove_utm'],
93     )
94
95     if serializer.data['handler'] == 'redirect':
96         return redirect(result['link'])
97
98     else:
99         return render(
100             request,
101             'bigly/link_info.html',
102             result,
103         )
104
105 class FollowRedirectsViewSet(viewsets.ViewSet):
106     serializer_class = serializers.FollowRedirectsSerializer
107
108     def follow_redirects(self, request):
109         serializer = serializers.FollowRedirectsSerializer(data=request.query_params)
110
111         if not serializer.is_valid():
112             return Response(
113                 serializer.errors,
114                 status=status.HTTP_400_BAD_REQUEST,
115             )
116
117         result = _follow_redirects(
118             link = serializer.data['link'],
119             remove_utm = serializer.data['remove_utm'],
120         )
121
122         return Response(
123             result,
124             status=status.HTTP_200_OK,
125         )
126
127 api_follow_redirects = FollowRedirectsViewSet.as_view({
128     'get': 'follow_redirects',
129 })