#!/usr/bin/env python

"""
Pure Python replacement for Apache's htpasswd

Borrowed from: https://gist.github.com/eculver/1420227

Modifications by James Mills, prologic at shortcircuit dot net dot au

- Added support for MD5 and SHA1 hashing.
"""


# Original author: Eli Carter

import os
import random
import sys
from hashlib import md5, sha1
from optparse import OptionParser


try:
    from crypt import crypt
except ImportError:
    try:
        from fcrypt import crypt
    except ImportError:
        crypt = None


def salt():
    """Returns a string of 2 randome letters"""
    letters = 'abcdefghijklmnopqrstuvwxyz' 'ABCDEFGHIJKLMNOPQRSTUVWXYZ' '0123456789/.'
    return random.choice(letters) + random.choice(letters)


class HtpasswdFile:
    """A class for manipulating htpasswd files."""

    def __init__(self, filename, create=False, encryption=None):
        self.filename = filename

        if encryption is None:
            self.encryption = lambda p: md5(p).hexdigest()
        else:
            self.encryption = encryption

        self.entries = []

        if not create:
            if os.path.exists(self.filename):
                self.load()
            else:
                raise Exception('%s does not exist' % self.filename)

    def load(self):
        """Read the htpasswd file into memory."""
        lines = open(self.filename).readlines()
        self.entries = []
        for line in lines:
            username, pwhash = line.split(':')
            entry = [username, pwhash.rstrip()]
            self.entries.append(entry)

    def save(self):
        """Write the htpasswd file to disk"""
        open(self.filename, 'w').writelines([f'{entry[0]}:{entry[1]}\n' for entry in self.entries])

    def update(self, username, password):
        """Replace the entry for the given user, or add it if new."""
        pwhash = self.encryption(password)
        matching_entries = [entry for entry in self.entries if entry[0] == username]
        if matching_entries:
            matching_entries[0][1] = pwhash
        else:
            self.entries.append([username, pwhash])

    def delete(self, username):
        """Remove the entry for the given user."""
        self.entries = [entry for entry in self.entries if entry[0] != username]


def main():
    """
    %prog [-c] -b filename username password
    Create or update an htpasswd file
    """
    # For now, we only care about the use cases that affect tests/functional.py
    parser = OptionParser(usage=main.__doc__)
    parser.add_option(
        '-b',
        action='store_true',
        dest='batch',
        default=False,
        help='Batch mode; password is passed on the command line IN THE CLEAR.',
    )
    parser.add_option(
        '-c',
        action='store_true',
        dest='create',
        default=False,
        help='Create a new htpasswd file, overwriting any existing file.',
    )
    parser.add_option(
        '-D', action='store_true', dest='delete_user', default=False, help='Remove the given user from the password file.'
    )

    if crypt is not None:
        parser.add_option('-d', action='store_true', dest='crypt', default=False, help='Use crypt() encryption for passwords.')

    parser.add_option('-m', action='store_true', dest='md5', default=False, help='Use MD5 encryption for passwords. (Default)')

    parser.add_option('-s', action='store_true', dest='sha', default=False, help='Use SHA encryption for passwords.')

    options, args = parser.parse_args()

    def syntax_error(msg):
        """
        Utility function for displaying fatal error messages with usage
        help.
        """
        sys.stderr.write('Syntax error: ' + msg)
        sys.stderr.write(parser.get_usage())
        sys.exit(1)

    if not options.batch:
        syntax_error('Only batch mode is supported\n')

    # Non-option arguments
    if len(args) < 2:
        syntax_error('Insufficient number of arguments.\n')
    filename, username = args[:2]
    if options.delete_user:
        if len(args) != 2:
            syntax_error('Incorrect number of arguments.\n')
        password = None
    else:
        if len(args) != 3:
            syntax_error('Incorrect number of arguments.\n')
        password = args[2]

    if options.crypt:

        def encryption(p):
            return crypt(p, salt())
    elif options.md5:

        def encryption(p):
            return md5(p).hexdigest()
    elif options.sha:

        def encryption(p):
            return sha1(p).hexdigest()
    else:

        def encryption(p):
            return md5(p).hexdigest()

    passwdfile = HtpasswdFile(filename, create=options.create, encryption=encryption)

    if options.delete_user:
        passwdfile.delete(username)
    else:
        passwdfile.update(username, password)

    passwdfile.save()


if __name__ == '__main__':
    main()
