#!/usr/bin/env python
# Based on FreeBSD src/lib/libcrypt/crypt.c 1.2
# http://www.freebsd.org/cgi/cvsweb.cgi/~checkout~/src/lib/libcrypt/crypt.c?rev=1.2&content-type=text/plain
#
# Original license:
# * "THE BEER-WARE LICENSE" (Revision 42):
# * <phk@login.dknet.dk> wrote this file.  As long as you retain this notice you
# * can do whatever you want with this stuff. If we meet some day, and you think
# * this stuff is worth it, you can buy me a beer in return.   Poul-Henning Kamp
#
# This port adds no further stipulations.  I forfeit any copyright interest.

from __future__ import print_function
import md5
import random
import string
import sys
import getpass

def hash(password, salt, magic='$apr1$'):
    # /* The password first, since that is what is most unknown */ /* Then our magic string */ /* Then the raw salt */
    m = md5.new()
    m.update(password + magic + salt)

    # /* Then just as many characters of the MD5(pw,salt,pw) */
    mixin = md5.md5(password + salt + password).digest()
    for i in range(0, len(password)):
        m.update(mixin[i % 16])

    # /* Then something really weird... */
    # Also really broken, as far as I can tell.  -m
    i = len(password)
    while i:
        if i & 1:
            m.update('\x00')
        else:
            m.update(password[0])
        i >>= 1

    final = m.digest()

    # /* and now, just to make sure things don't run too fast */
    for i in range(1000):
        m2 = md5.md5()
        if i & 1:
            m2.update(password)
        else:
            m2.update(final)

        if i % 3:
            m2.update(salt)

        if i % 7:
            m2.update(password)

        if i & 1:
            m2.update(final)
        else:
            m2.update(password)

        final = m2.digest()

    # This is the bit that uses to64() in the original code.

    itoa64 = './0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'

    rearranged = ''
    for a, b, c in ((0, 6, 12), (1, 7, 13), (2, 8, 14), (3, 9, 15), (4, 10, 5)):
        v = ord(final[a]) << 16 | ord(final[b]) << 8 | ord(final[c])
        for i in range(4):
            rearranged += itoa64[v & 0x3f]; v >>= 6

    v = ord(final[11])
    for i in range(2):
        rearranged += itoa64[v & 0x3f]; v >>= 6

    return magic + salt + '$' + rearranged

def usage():
    print('%s: usage: %s [-cD] passwdfile username' %
        (sys.argv[0], sys.argv[0]), file=sys.stderr)

def salt(len):
    return ''.join([random.choice(string.ascii_letters + string.digits)
        for x in range(len)])

def write_passwords(passwords, path):
    with open(path, 'w') as f:
        for (username, pwhash) in passwords:
            f.write('%s:%s\n' % (username, pwhash))

def ask_password():
    x = getpass.getpass('New password: ')
    y = getpass.getpass('Re-type password: ')

    return x if x == y else None

if len(sys.argv) not in [3, 4]:
    if len(sys.argv) == 4 and sys.argv[1] not in ['-c', '-D', '-cD', '-Dc']:
        usage()
        sys.exit(2)

flags = len(sys.argv) == 4
create_flag = flags and 'c' in sys.argv[1]
delete_flag = flags and 'D' in sys.argv[1]

if create_flag and delete_flag:
    print('%s: -c and -D options conflict' % sys.argv[0], file=sys.stderr)
    sys.exit(2)

file_path = sys.argv[flags + 1]
username = sys.argv[flags + 2]

if not delete_flag:
    password = ask_password()

    if password == None:
        exit("%s: passwords weren't the same" % sys.argv[0])

contents = []
found = False

if not create_flag:
    with open(file_path, 'r') as f:
        # read in the existing passwd file
        # replace entry for 'username' with entry containing new hash
        # unless -D is used, in which case we remove the entry
        #
        # example entry: username:$apr1$gdehCd2T$ppFjRXlf1alPKSHqcBrjk0

        for line in f:
            (u, ph) = string.split(line.strip('\n'), ':')

            if u == username:
                if not delete_flag:
                    ph = hash(password, salt(8))
                    print('Updating password for user %s' % username)
                    contents.append((u, ph))

                found = True
            else:
                contents.append((u, ph))

if not found:
    if delete_flag:
        print('User %s not found' % username)
    else:
        print('Adding password for user %s' % username)
        contents.append((username, hash(password, salt(8))))

write_passwords(contents, file_path)