shortener/shortener/resolver.py

import os

from dns.name import from_text
from dns.resolver import Resolver

from shortener.exceptions import ShortCodeNotFoundException

BASE_DOMAIN = os.getenv("BASE_DOMAIN")


class ShortenerResolver(Resolver):
    def __init__(self):
        super().__init__()
        self.search = [from_text(BASE_DOMAIN + ".")]

    def resolve_shortcode(self, shortcode):
        try:
            answer = self.resolve(shortcode, "TXT", search=True, lifetime=10)
            redirect_url = list(answer)[0].strings[0].decode("utf-8")
            return redirect_url
        except:
            raise ShortCodeNotFoundException()


resolver = ShortenerResolver()