#!/usr/libexec/platform-python -E

#
# lookup CERN LDAP for printer info and setup local CUPS spooler with right printer options.
#
# May 2022 - Ben Morrice <ben.morrice@cern.ch>
# - port from perl to python
# Mar 2010 - Feb 2014 Jaroslaw.Polok <jaroslaw.polok@cern.ch>

import argparse
import os
import re
import subprocess
import sys
import ldap

LDAP_SERVER = 'xldap.cern.ch'


def ldap_init():
    try:
        ldap.set_option(ldap.OPT_NETWORK_TIMEOUT, 5.0)
        ldap.set_option(ldap.OPT_TIMEOUT, 10)
        conn = ldap.initialize(f"ldap://{LDAP_SERVER}:389")
        conn.simple_bind_s()
    except (ldap.LDAPError,ldap.SERVER_DOWN) as error_message:
        print(f"Cannot connect to ldap. Error: {error_message}")
        sys.exit(50)
    return conn


def ldap_search(conn, base, search, attrs, verbose):
    if verbose:
        print(f"LDAP search: {search}")
    result = conn.search_s(base, ldap.SCOPE_SUBTREE, search, attrs)
    if verbose:
        print(f"LDAP return: {len(result)} entries.")
    if len(result) == 0:
        print('No match found.')
        sys.exit(10)
    return result


def add_printer(printer, verbose):
    output = run_cmd(f"/usr/bin/lpstat -p '{printer[0]}'", verbose)[1]
    test = re.search(f"printer {printer[0]}", output.decode())
    if test is not None:
        print(f"Updating printer {printer[0]}")
    else:
        print(f"Adding printer {printer[0]}")
    if printer[1]['serverName'][0]:
        uri = f"lpd://{printer[1]['serverName']}/{printer[0]}"
    else:
        uri = f"ldp://{printer[0]}"
    if 'TRUE' in printer[1]['printDuplexSupported']:
        duplex = '-o Duplex=DuplexNoTumble -o Duplexer=True -o DuplexUnit=Installed -o HPOption_Duplexer=True -o OKOptionDuplex=True -o Option1=True -o Option3=True -o XRXOptionDuplex=True -o DefaultDuplex=DuplexNoTumble'
    else:
        duplex = ''
    run_cmd(
        f"/usr/sbin/lpadmin -p {printer[0]} -L \"{printer[1]['location']}\" -P {printer[1]['ppdfile']} -D \"{printer[1]['description']}\" -v {uri} {duplex} -o PageSize=A4 -o printer-is-shared=false -E", verbose)


def remove_printer(printer, verbose):
    if run_cmd(f"/usr/bin/lpstat -p '{printer}'", verbose)[0] == 0:
        print(f"Removing printer: {printer}")
        run_cmd(f"/usr/sbin/lpadmin -x '{printer}'", verbose)[0]


def get_local_printers(printers):
    local_printers = {}
    try:
        with open('/etc/cups/printers.conf', 'r', encoding='unicode-escape') as f:
            for line in f:
                if '<Printer' in line:
                    printer = re.search('<Printer (.*)>', line).group(1)
                    for k, v in printers.items():
                        if k in printer:
                            local_printers.update({k: v})
    # Only root can read the above file, but we don't need it for 'list' purposes
    except PermissionError:
        pass
    return local_printers


def get_default_printer(verbose):
    output = run_cmd('/usr/bin/lpstat -d', verbose)[1]
    test = re.search(r'^system\sdefault\sdestination.\s+(.*)', output.decode())
    if test:
        if verbose:
            print(f"System default printer: {test.group(1)}")
        return test.group(1)
    return None


def set_default_printer(default, verbose):
    if default is None:
        return
    if run_cmd(f"/usr/bin/lpstat -p '{default}'", verbose)[0] != 0:
        print(
            f"Error: cannot set default printer to {default} (no such printer)")
        sys.exit(1)

    print(f"Setting default printer to: {default} ...")
    run_cmd(f"/usr/sbin/lpadmin -d '{default}'", verbose)


def get_print_options(printer, verbose):
    if run_cmd(f"/usr/bin/lpstat -p '{printer}'", verbose)[0] != 0:
        return None
    output = run_cmd(f"/usr/bin/lpoptions -p '{printer}' -l", verbose)
    options = ''
    for item in output[1].splitlines():
        try:
            keys, values = item.decode().split(':')
        except ValueError:
            continue
        if not '/' in keys:
            continue
        key = keys.split('/')[0]
        value = re.search(r'\*(\w+)', values).group(1)
        if not value:
            continue
        value.replace('*', '')
        options += f" -o {key}={value}"
    return options


def set_print_options(printer, options, verbose):
    if options is not None:
        run_cmd(f"/usr/bin/lpoptions -p {printer} {options}", verbose)


def reloadcupsd(verbose):
    run_cmd('/usr/bin/systemctl restart cups', verbose)


def run_cmd(cmd, verbose):
    cmd = f"LC_ALL=C {cmd}"
    if verbose:
        print(f"running: {cmd}")
    try:
        p = subprocess.Popen(cmd, stderr=subprocess.PIPE,
                             shell=True, stdout=subprocess.PIPE)
    except:
        if verbose:
            print(f"Failed to execute: {cmd}")
        return 1
    out, err = p.communicate()
    if p.wait() != 0:
        if verbose:
            print(f"Error executing command: {cmd}")
    return p.wait(), out, err


def printer_driver_bindings(verbose):
    nametweaks = {
        # LDAP returns from landb Model in PPD file
        "CANON IR1730": "Canon iR1730i PS3",
        "CANON IR-ADV C5030": "Canon iR-ADV C5030/5035 PS",
        "CANON IR-ADV 4225": "Canon iR-ADV 4225/4235 PS",
        "CANON IR C3080I": "Canon iR C3080/3480/3580 PS",
        "CANON IR C3080": "Canon iR C3080/3480/3580 PS",
        "CANON IR2270": "Canon iR2270/iR2870 PS",
        "CANON IR C 2380 I": "Canon iR C2380/2550 PS",
        "CANON IR C2380": "Canon iR C2380/2550 PS",
        "CANON IR1022A": "Canon iR1018-1022 CPCA",
        "CANON IR1024I": "Canon iR1018-1022 CPCA",  # VERIFY
        "CANON IR1024": "Canon iR1018-1022 CPCA",  # VERIFY
        "CANON IR3035": "Canon iR3035/iR3045 PS",
        "CANON IR3045": "Canon iR3035/iR3045 PS",
        "CANON IRC 3380I": "Canon iR C2880/C3380 PS",
        "CANON IR C3380": "Canon iR C2880/C3380 PS",
        "CANON IRC 2880 I": "Canon iR C2880/C3380 PS",
        "CANON IR C2880": "Canon iR C2880/C3380 PS",
        "CANON IR3300I": "Canon iR2200-3300 PS",
        "CANON IR-ADV 4245": "Canon iR-ADV 4245/4251 PS",
        "CANON IR-ADV 4245": "Canon iR-ADV 4245/4251 PS",
        "CANON IR4570": "Canon iR3570/iR4570 PS",
        "CANON IR3570": "Canon iR3570/iR4570 PS",
        "CANON IR3245": "Canon iR3235/iR3245 PS",
        "CANON IR 3245N": "Canon iR3235/iR3245 PS",
        "CANON IR3235N/E": "Canon iR3235/iR3245 PS",
        "CANON IR3235": "Canon iR3235/iR3245 PS",
        "CANON IR 5055N": "Canon iR5055/iR5065 PS",
        "CANON IR5055": "Canon iR5055/iR5065 PS",
        "CANON IR 6570": "Canon iR5570/iR6570 PS",
        "CANON IR6570": "Canon iR5570/iR6570 PS",
        "CANON IR 5870": "Canon iR C5870C EUR PS",
        "CANON IR 5870C": "Canon iR C5870C EUR PS",
        "CANON IR 2570C EUR": "Canon iR 2570C EUR PS",
        "CANON IR 2570C": "Canon iR 2570C EUR PS",
        "CANON LBP6650": "Canon LBP6650/3470 PS",
        "CANON IR-ADV C5535": "Canon iR-ADV C5535/5540 PS",
        "CANON IR-ADV C5235": "Canon iR-ADV C5235/5240 PS",
        "CANON IR-ADV 4025": "Canon iR-ADV 4025/4035 PS",
        "CANON IR-ADV 4045": "Canon iR-ADV 4045/4051 PS",
        "CANON IR3225": "Canon iR3225 PS",
        "CANON IR-ADV 400": "Canon iR-ADV 400/500 PS",
        "CANON IR-ADV C5535I": "Canon iR-ADV C5535/5540 PS",
        "CANON IR-ADV C356I": "Canon iR-ADV C256/356 UFR II",
        "CANON IR-ADV C356": "Canon iR-ADV C256/356 UFR II",
        "CANON IR-ADV 4535": "Canon iR-ADV 4525/4535 PS",
        "CANON IR-ADV 525": "Canon iR-ADV 525 PS",
        "CANON IR-ADV 4735": "Canon iR-ADV 4725/4735 PS",
        "CANON IR-ADV 4935": "Canon iR-ADV 4935 PS",
        "CANON IR-ADV 527": "Canon iR-ADV DX 527 PS",
        "CANON IR-ADV 529": "Canon iR-ADV 529 PS",
        "CANON IR-ADV C357": "Canon iR-ADV DX C257/357 PS",
        "CANON IR-ADV C359": "Canon iR-ADV C259/359 PS",
        "CANON IR-ADV C5735": "Canon iR-ADV C5735/5740 PS",
        "CANON IR-ADV C5740": "Canon iR-ADV C5735/5740 PS",
        "CANON IR-ADV C5840": "Canon iR-ADV DX C5840/5850 PS",
        "HP LASERJET 500 COLOR M551": "HP LaserJet 500 color M551",
        "HP COLOR LASERJET M553": "HP Color LaserJet M553",
        "HP COLOR LASERJET CP5520 SERIES": "HP Color LaserJet CP5520 Series",
        "HP LASERJET 600 M603": "HP LaserJet 600 M601 M602 M603",
        "HP Laserjet P2055DN": "HP LaserJet P2055 Series",
        "HP LASERJET P2055DN": "HP LaserJet P2055 Series",
        "HP LASERJET 2200": "HP LaserJet 2200",
        "HP LASERJET P2015 SERIES": "HP LaserJet P2015 Series",
        "HP LASERJET P3004": "HP LaserJet P3004",
        "HP LASERJET P4014": "HP LaserJet P4010 Series",
        "HP LASERJET 4300": "HP LaserJet 4300 Series",
        "HP LASERJET 4100 SERIES": "HP LaserJet 4100 Series ",  # trailing space !
        "HP LASERJET P3010 SERIES": "HP LaserJet P3010 Series",
        "HP LASERJET 2430PCL6": "HP LaserJet 2430",
        "HP LASERJET 2300 DTN": "HP LaserJet 2300",
        "HP LASERJET 2300 SERIES": "HP LaserJet 2300",
        "HP COLOR LASERJET CP 6015 XH": "HP Color LaserJet CP6015",
        "HP COLOR LASERJET CM4730 MFP": "HP Color LaserJet 4730mfp",
        "HP BUSINESS INKJET 2800": "HP Business Inkjet 2800 PS",
        "HP BUISNESS INKJET 2800": "HP Business Inkjet 2800 PS",  # (sp.!)
        "HP LASERJET IIISI POSTSCRIPT": "HP LaserJet 3",
        "HP LASERJET 5M": "HP LaserJet 5/5M",
        "HP LASERJET 5SI": "HP LaserJet 5Si/5Si MX",
        "HP LASERJET 6MP": "HP LaserJet 6P/6MP",
        "HP LASERJET P3015DN": "HP LaserJet 3015",
        "HP LASERJET 2100 TN": "HP LaserJet 2100 Series",
        "HP LASERJET 2200 SERIES": "HP LaserJet 2200",
        "HP LASERJET 4050 SERIES": "HP LaserJet 4000 Series",
        "HP LASERJET 4250": "HP LaserJet 4200",
        "HP LASERJET P4515": "HP LaserJet P4010 Series",
        "HP LASERJET 700 M712": "HP LaserJet 700 M712",
        "HP COLOR LASERJET 2840": "HP Color LaserJet 2800",
        "HP LASERJET 5100 SERIES": "HP LaserJet 5100",
        "HP LASERJET 8150 SERIES": "HP LaserJet 8150 Series PS",
        "HP HP COLOR LASERJET 4550": "HP Color LaserJet 4550 ",  # trailing space !
        "HP COLOR LASERJET 5500": "HP Color LaserJet 5500 PS",
        "HP 2500C SERIES": "HP DesignJet 2500CP PS3",
        "HP DESIGNJET 2500CP": "HP DesignJet 2500CP PS3",
        "HP DESIGNJET 650C .C2859A.": "HP DesignJet 650C",
        "HP DESIGNJET 650C .C2859B.": "HP DesignJet 650C",
        "HP DESIGNJET 750C": "HP DesignJet 750C Plus",
        "HP DESIGNJET 755C": "HP DesignJet 750C Plus",
        "HP DESIGNJET 755CM": "HP DesignJet 750C Plus",
        "HP DESIGNJET 1055CM": "HP DesignJet 1055CM PS3",
        "HP LASERJET M606": "HP Laserjet M604 M605 M606",
        "HP LASERJET M506": "HP LaserJet M506",
        "INFOTEC ISC 2428": "infotec  ISC 2428",
        "OCE 3165": "Oce 3165 PS3",
        "OCE TDS450": "Oce TDS450 PS",
        "OCE TCS500 COLOUR": "Oce TCS500 PS",
        "OCE TCS500": "Oce TCS500 PS",
        "OCE TDS600": "Oce TDS600 PS",
        "OCE TDS700": "Oce TDS700 PS",
        "TEKTRONIX PHASER 750 DX": "Tektronix Phaser 750DP",
        "TEKTRONIX PHASER 750DX": "Tektronix Phaser 750DP",
        "TEKTRONIX PHASER 750": "Tektronix Phaser 750DP",
        "TEKTRONIX PHASER 850DP": "Tektronix Phaser 860DP",
        "TEKTRONIX PHASER 8200DP": "Xerox Phaser 8200DP",
        "XEROX DOCUMENT CENTRE 440": "Xerox Document Centre 440 PS",
        "XEROX DOCUMENT CENTRE 535": "Xerox Generic",  # Xerox could really try harder
        "XEROX WORKCENTRE 255": "Xerox Generic",  # Xerox could really try harder
        "XEROX WORKCENTRE PRO 255": "Xerox Generic",  # Xerox could really try harder
        "XEROX WORKCENTRE 245 PRO": "Xerox Generic",  # Xerox could really try harder
        "XEROX WORKCENTRE PRO 245": "Xerox Generic",  # Xerox could really try harder
        "XEROX WORKCENTRE 7245": "Xerox Generic",  # Xerox could really try harder
        "XEROX PHASER 860": "Xerox Generic",  # Xerox could really try harder
        "XEROX PHASER 7300 DN": "Xerox Phaser 7300DN",
        "XEROX PHASER 4510N": "Xerox Generic",  # Xerox could really try harder
        "XEROX COLORQUBE 9203": "Xerox ColorQube 9201/9202/9203"
    }

    printer_models = {}
    ppd_path = '/usr/share/lpadmincern/ppds/'
    for ppd in os.listdir(ppd_path):
        with open(os.path.join(ppd_path, ppd), 'r', encoding='unicode-escape') as f:
            for line in f:
                if isinstance(line, str):
                    if re.search(r'\*ModelName', line):
                        printer_model = line.split(
                            ':')[1].replace('"', '').strip()
                        if verbose:
                            print(f"Printer model: {printer_model} PPD: {ppd}")
                        printer_models[printer_model] = f"{ppd_path}{ppd}"

    # tweaks
    for k, v in nametweaks.items():
        try:
            printer_models[k] = printer_models[v]
        except KeyError:
            pass
    return printer_models


def main():

    parser = argparse.ArgumentParser()
    functional_group = parser.add_mutually_exclusive_group()
    color_group = parser.add_mutually_exclusive_group()
    duplex_group = parser.add_mutually_exclusive_group()
    parser.add_argument('--name', required=False, type=str,
                        action='store', help='Search for printers matching PRINTERNAME')
    parser.add_argument('--model', required=False, type=str,
                        action='store', help='Search for printers of MODEL')
    parser.add_argument('--building', required=False, type=str, action='append',
                        nargs='+', metavar={}, help='Search for printers in BUILDING(s)')
    color_group.add_argument('--color', required=False,
                             action='store_true', help='Search for Color printers only')
    color_group.add_argument('--nocolor', required=False,
                             action='store_true', help='Search for monochrome printers only')
    duplex_group.add_argument('--duplex', required=False, action='store_true',
                              help='Search for DUPLEX capable printers only')
    duplex_group.add_argument('--noduplex', required=False, action='store_true',
                              help='Search for non DUPLEX capable printers only')
    parser.add_argument('--paper', required=False, type=str,
                        action='store', help='Search for printers allowing PAPERSIZE')
    functional_group.add_argument(
        '--list', required=False, action='store_true', help='Show matching printers')
    functional_group.add_argument(
        '--add', required=False, action='store_true', help='Add matching printers to the system')
    functional_group.add_argument(
        '--remove', required=False, action='store_true', help='Remove matching printers from the system')
    functional_group.add_argument('--update', required=False, action='store_true',
                                  help='Update matching printer definitions on the system')
    parser.add_argument('--default', required=False, type=str,
                        action='store', help='Set default system printer to PRINTER')
    parser.add_argument('--verbose', required=False,
                        action='store_true', help=argparse.SUPPRESS)

    # This logic ensures that single dash shortname arguments work (eg: '-build 31')
    new_argv = []
    for arg in sys.argv:
        if arg.startswith('-') and not arg.startswith('--') and len(arg) > 2:
            arg = '-' + str(arg)
        new_argv.append(str(arg))

    sys.argv = new_argv
    # This logic ensures that a user can pass [NAME] without any other arguments
    # (the same behaviour as the perl version)
    args, unknown = parser.parse_known_args()
    if not args.name and len(unknown) == 1:
        args.name = unknown[0]
    verbose = args.verbose

    search_string = '(&'
    if args.update or args.default:
        search_string += '(printerName=*)'
    if args.name:
        search_string += f"(printerName={args.name})"
    if args.model:
        search_string += f"(description={args.model})"
    if args.color:
        search_string += '(printColor=TRUE)'
    if args.nocolor:
        search_string += '(printColor=FALSE)'
    if args.duplex:
        search_string += '(printDuplexSupported=TRUE)'
    if args.noduplex:
        search_string += '(printDuplexSupported=FALSE)'
    if args.paper:
        search_string += f"(printMediaSupported=*{args.paper}*)"
    if args.building:
        search_string += '(|'
        for b in args.building:
            search_string += f"(location={b[0]} *)"
        search_string += ')'
    search_string += '(objectClass=printQueue)'
    search_string += ')'

    conn = ldap_init()
    printer_results = ldap_search(conn, 'ou=print servers,ou=system infrastructure,ou=cern main servers,dc=cern,dc=ch', search_string, [
                                  'printerName', 'location', 'description', 'printColor', 'printDuplexSupported', 'serverName', 'portName', 'printMediaReady'], verbose)

    for printer in range(0, len(printer_results)):
        for k, v in printer_results[printer][1].items():
            if len(v) > 1:
                printer_results[printer][1][k] = ''
                for media in v:
                    printer_results[printer][1][k] = f"{printer_results[printer][1][k]} {media.decode()}".strip()
            else:
                printer_results[printer][1][k] = v[0].decode()

    printerbindings = printer_driver_bindings(verbose)
    allprinters = {}
    for p in printer_results:
        if 'printerName' not in p[1]:
          continue
        try:
          printer = {p[1]['printerName'].upper(): {'description': p[1]['description'], 'driver': 'GENERIC POSTSCRIPT PRINTER', 'ppdfile': printerbindings['Generic PostScript Printer'],
                                                 'location': p[1]['location'].strip(), 'printColor': p[1]['printColor'], 'printDuplexSupported': p[1]['printDuplexSupported'], 'serverName': p[1]['serverName']}}
        except KeyError:
            print(f"Rejected data for: {p[1]['printerName'].upper()} (incomplete LDAP object).")
            continue
        for k in printerbindings.items():
            if p[1]['description'].upper() in k[0].upper():
                printer[p[1]['printerName'].upper()]['driver'] = k[0].upper()
                printer[p[1]['printerName'].upper()]['ppdfile'] = k[1]
        allprinters.update(printer)
        if verbose:
            print(f"Binding: {p[1]['printerName'].upper()} = {p[1]['description']} -> {printer[p[1]['printerName'].upper()]['driver']} [{printer[p[1]['printerName'].upper()]['ppdfile']}]")

    localprinters = get_local_printers(allprinters)

    # Ensure that --list is the default option
    if not args.add and not args.remove and not args.update and not args.default:
        args.list = True

    if not args.list:
        if os.geteuid() != 0:
            print('You must be root in order to add/remove/update/set default printer(s)')
            sys.exit(1)

    if args.list:
        print(f"{'Printer name' :<22} {'Printer model' :<31} {'Location' :<11} {'Color':<5} {'Duplex' :<3}")
        print('-'*80)
        for p, v in sorted(allprinters.items()):
            if 'TRUE' in v['printColor']:
                printColor = 'yes'
            else:
                printColor = 'no'
            if 'TRUE' in v['printDuplexSupported']:
                printDuplexSupported = 'yes'
            else:
                printDuplexSupported = 'no'
            if 'GENERIC POSTSCRIPT PRINTER' in v['driver']:
                genericDriver = ' !'
            else:
                genericDriver = ''
            print(
                f"{p :<22} {v['description'] :<31} {v['location'] :<11} {printColor :<5} {printDuplexSupported :<3}{genericDriver}")
        print('-'*80)
        print(
            f"{len(allprinters)} matches found. (! in last column: generic postscript driver used).")
        sys.exit(0)

    if args.add:
        for p in allprinters.items():
            add_printer(p, verbose)

    if args.remove:
        for p in allprinters:
            remove_printer(p, verbose)

    if args.update:
        default = get_default_printer(verbose)
        for p in localprinters.items():
            options = get_print_options(p[0], verbose)
            add_printer(p, verbose)
            set_print_options(p[0], options, verbose)
        if default is not None:
            set_default_printer(default, verbose)

    if args.default:
        set_default_printer(args.default, verbose)

    # reload cups in all cases except for listing
    reloadcupsd(verbose)


if __name__ == "__main__":
    main()
