# Copyright 2025 Google LLC

import random
import socket
import string
import datetime

class Banner(object):
    """Act like a string, but actually get date/time components on the fly.

    Returned by BannerFactory.genBanner().
    
    Allows listeners to statically set a banner in their start() procedure for
    libraries like pyftpdlib that expect to be able to use a static string
    value. When the __len__() and __repr__() methods are called, this class
    dynamically uses datetime.datetime.now().strftime() and the provided
    insertion strings to ensure that the banner string reflects the current
    time and date at the time when it was referenced, which allows support for
    banners like those used by wu-ftpd and others that reference the current
    date, time, servername, etc.
    """

    def __init__(self, banner, insertions):
        # Use this to enable/disable test path to simulate length of 75 but
        # actual string of length 76 to test what happens around
        # pyftpdlib/handlers.py:1321 when the return value of len() is <= the
        # threshold of 75 characters but the string's actual length (which is
        # generated in a subsequent call to __repr__()) exceeds that threshold
        # value. The conclusion as of this writing is that a discrepancy in the
        # length returned by the __len__() method and the actual length of the
        # string returned by __repr__() does not affect FTP functionality, so
        # the corner case in which these may differ (due to the generated
        # string being longer or shorter after it is evaluated due to elements
        # interpolated by strftime) will be ignored.
        self.test_pyftpdlib_handler_banner_threshold75 = False

        if self.test_pyftpdlib_handler_banner_threshold75:
            self.len_75 = len(self.str_75)
            self.str_76 = 'a' * 76

        self.banner = banner
        self.insertions = insertions

        # Indicate an error in the banner early-on as opposed to
        # when a login or other event occurs.
        testbanner, testlen = self.failEarly()

    def failEarly(self):
        """Raise exceptions upon construction rather than later."""

        # Test generating banner
        banner_generated = str(self)

        # Test generating and getting length of banner
        banner_generated_len = len(self)

        return banner_generated, banner_generated_len

    def __len__(self):
        """Needed for pyftpdlib.
        
        If the length changes between the time when the caller obtains the
        length and the time when the caller obtains the latest generated
        string, then there is not much that could reasonably be done. It would
        be possible to cache the formatted banner with a short expiry so that
        temporally clustered __len__() and __repr__() call sequences would view
        consistent and coherent string contents, however this seems like
        overkill since the use case is really just allowing pyftpdlib to
        determine which way to send the response (directly versus push() if the
        length exceeds a threshold of 75 characters). In this case, if the
        banner string length and contents are inconsistent, it appears that the
        only effect will be to erroneously send the message differently. Test
        code has been left in place for easy repro in case this proves to be an
        issue on some future/other platform.
        """
        # Test path: simulate length of 75 but actual string of length 76 (part
        # 1/2) to test pyftpdlib/handlers.py:1321
        if self.test_pyftpdlib_handler_banner_threshold75:
            return self.len_75

        # Normal path: return the length of the banner generated by self.fmt()
        return len(self.fmt())

    def __repr__(self):
        return self.fmt()

    def fmt(self):
        # Test path: simulate length of 75 but actual string of length 76 (part
        # 2/2) to test pyftpdlib/handlers.py:1321
        if self.test_pyftpdlib_handler_banner_threshold75:
            return self.str_76

        # Normal path: generate banner
        banner = self.banner
        banner = banner.format(**self.insertions)
        banner = datetime.datetime.now().strftime(banner)
        banner = banner.replace('\\n', '\n').replace('\\t', '\t')
        return banner


class BannerFactory(object):
    def genBanner(self, config, bannerdict, defaultbannerkey='!generic'):
        """Select and initialize a banner.
        
        Supported banner escapes:
            !<key> - Use the banner whose key in bannerdict is <key>
            !random - Use a random banner from bannerdict
            !generic - Every listener supporting banners must have a generic

        Banners can include literal '\n' or '\t' tokens (slash followed by the
        letter n or t) to indicate that a newline or tab should be inserted.

        Banners can include {servername} or {tz} to insert the servername or
        time zone (hard-coded to 'UTC' as of this writing).

        If the user does not specify a banner, then '!generic' is used by
        default, resulting in bannerdict['generic'] being used. If the user
        specifies a bang escape e.g. '!iis-6', then the banner keyed by that
        name will be used. If the user specifies '!random' then a random banner
        will be chosen from bannerdict.

        Because some banners include the servername as an insertion string,
        this method also retrieves the configuration value for ServerName and
        incorporates a couple of similar escape sequences:
            !random - Randomized servername with random length between 1-15
            !gethostname - Use the real hostname
        """

        banner = config.get('banner', defaultbannerkey)
        servername = config.get('servername', 'localhost')

        if servername.startswith('!'):
            servername = servername[1:]
            if servername.lower() == 'random':
                servername = self.randomizeHostname()
            elif servername.lower() == 'gethostname':
                servername = socket.gethostname()
            else:
                raise ValueError('ServerName config invalid escape: !%s' %
                        (servername))

        if banner.startswith('!'):
            banner = banner[1:]
            if banner.lower() == 'random':
                banner = random.choice(list(bannerdict.keys()))
            elif banner not in bannerdict:
                raise ValueError(
                        'Banner config escape !%s not a valid banner key' %
                        (banner))

            banner = bannerdict[banner]

        insertions = {'servername': servername, 'tz': 'UTC'}

        return Banner(banner, insertions)

    def randomizeHostname(self):
        valid_hostname_charset = (string.ascii_letters + string.digits + '-')
        n = random.randint(1, 15)
        return ''.join(random.choice(valid_hostname_charset) for _ in range(n))
