sharp-postprocess

#! /usr/bin/env python

from Crypto.Cipher import AES
import base64
import sys
import os
import time
import re
import fcntl
import select
import pwd
import io

# ****************************************

def ask_user(title, questions, username):
    import dbus

    user = pwd.getpwnam(username)
    if user is None:
        return None
    uid = user[2]

    try:
        bus = dbus.SystemBus()

        bus_object = bus.get_object('org.freedesktop.DBus',
                                    '/org/freedesktop/DBus')
        bus_iface = dbus.Interface(bus_object, 'org.freedesktop.DBus')
        owners = bus_iface.ListQueuedOwners('org.cups.sharp.Passwd')


        for owner in owners:
            remote_uid = bus_iface.GetConnectionUnixUser(owner)
            # Only send request to the correct user.
            if remote_uid != uid:
                continue

            remote_object = bus.get_object(owner, '/org/cups/sharp/Passwd')

            iface = dbus.Interface(remote_object, "org.cups.sharp.Passwd")
            reply = iface.Ask(title, questions, timeout=3600)

            # We got a reply, so we are done!
            return reply
        sys.stderr.write("ERROR: User does not seem to have the service running.\n")
    except dbus.DBusException, e:
        sys.stderr.write("ERROR: Could not ask user for authentication data.\n")
        sys.stderr.write(str(e))
        return None

    return None

# ****************************************

def encrypt(passwd, spooltime):
    key = spooltime[-16:]
    # We need to add at least one NULL byte
    missing = 16 - len(passwd) % 16
    passwd += '\0'*missing

    aes = AES.new(key=key, mode=AES.MODE_ECB)
    encrypted = aes.encrypt(passwd)
    encrypted = base64.encodestring(encrypted).strip()

    return encrypted

jobid   = sys.argv[1]
user    = sys.argv[2]
jobname = sys.argv[3]
copies  = int(sys.argv[4])
options = sys.argv[5]
if len(sys.argv) > 6:
    filename=sys.argv[6]
    input_fd = open(filename)
else:
    input_fd = sys.stdin

output_fd = sys.stdout

spooltime=time.strftime("%a, %d %h %Y %H:%M:%S %z", time.gmtime())

option_regex = re.compile("^@PJL\s*COMMENT\s*XSHARP_(?P<option>.*?)\s*=\s*(?P<argument>.*)$")
aes_encrypt_regex = re.compile("XSHARP_AES_ENCRYPT:(?P<input>.*?):XSHARP_AES_ENCRYPT")
askpass_regex = re.compile("XSHARP_ASK_PASSWORD:(?P<label>.*?):XSHARP_ASK_PASSWORD")
askuser_regex = re.compile("XSHARP_ASK_STRING:(?P<label>.*?):XSHARP_ASK_STRING")
askpin_regex = re.compile("XSHARP_ASK_PASSCODE:(?P<label>.*?):XSHARP_ASK_PASSCODE")


# ****************************************
# Now read the file line by line, until we find the "@PJL ENTER" header.

def replace_argument(opt, arg, user_input=None):
    # OK, process the argument ...
    if user_input is None:
        input_needed = []

        for label in askpass_regex.findall(arg):
            input_needed.append((opt+'::'+label, label, 'PASSWORD'))
        for label in askuser_regex.findall(arg):
            input_needed.append((opt+'::'+label, label, 'STRING'))
        for label in askpin_regex.findall(arg):
            input_needed.append((opt+'::'+label, label, 'PASSCODE'))
        if len(input_needed) > 0:
            return None, input_needed

    arg = arg.replace("XSHARP_SPOOLTIME", spooltime)
    arg = arg.replace("XSHARP_USERNAME", user)
    arg = arg.replace("XSHARP_COPIES", str(copies))

    def askpass_replace(match):
        if user_input is not None and user_input.has_key(opt+'::'+match.group('label')):
            return user_input[opt+'::'+match.group('label')]
    arg, cnt = askpass_regex.subn(askpass_replace, arg)

    def askuser_replace(match):
        if user_input is not None and user_input.has_key(opt+'::'+match.group('label')):
            return user_input[opt+'::'+match.group('label')]
    arg, cnt = askuser_regex.subn(askuser_replace, arg)

    def askpin_replace(match):
        if user_input is not None and user_input.has_key(opt+'::'+match.group('label')):
            return user_input[opt+'::'+match.group('label')]
    arg, cnt = askpin_regex.subn(askpin_replace, arg)

    def regex_encrypt(match):
        return encrypt(match.group("input"), spooltime)
    arg, cnt = aes_encrypt_regex.subn(regex_encrypt, arg)

    return arg, None

buf = []
input_needed = None
while True:
    line = input_fd.readline()
    if not line:
        sys.exit(0)

    match = option_regex.match(line)
    if match:
        # Got a comment to process ...
        option = match.group("option")
        orig_arg = match.group("argument")

        arg, new_input_needed = replace_argument(option, orig_arg)
        if new_input_needed:
            if input_needed is None:
                input_needed = new_input_needed
            else:
                input_needed += new_input_needed

        if input_needed:
            buf.append((option, orig_arg))
        else:
            output_fd.write("@PJL SET %s=%s\n" % (option, arg))
        continue

    if line.startswith("@PJL ENTER"):
        # We may need to get input from the user ...
        if input_needed:
            user_input = ask_user(jobname, input_needed, user)
            # catch errors ...
            if user_input is None:
                user_input = {}
            for option, orig_arg in buf:
                if option is None:
                    output_fd.write(orig_arg)
                else:
                    arg, new_input_needed = replace_argument(option, orig_arg, user_input)
                    output_fd.write("@PJL SET %s=%s\n" % (option, arg))

        output_fd.write(line)
        break

    if input_needed:
        buf.append((None, line))
    else:
        output_fd.write(line)

# There appears to be a bug, the output buffer is not flushed
# when enabling non-blocking IO. And then python tries to
# write it *after* the fd is closed.
output_fd.flush()

in_fd = input_fd.fileno()
fl = fcntl.fcntl(in_fd, fcntl.F_GETFL)
fcntl.fcntl(in_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)

out_fd = output_fd.fileno()
fl = fcntl.fcntl(out_fd, fcntl.F_GETFL)
fcntl.fcntl(out_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)

# We need to use io.open so that the BlockingIOError is raised.
output_fd = io.open(out_fd, mode="bw", buffering=8192)

buf = []

out_open = True
in_open = True

while (out_open and len(buf) > 0) or (out_open and in_open):
    poll = select.poll()

    if out_open:
        if len(buf) > 0:
            mask = select.POLLOUT | select.POLLHUP | select.POLLERR
        else:
            mask = select.POLLHUP | select.POLLERR
        poll.register(out_fd, mask)

    if in_open:
        if len(buf) < 100:
            mask = select.POLLIN | select.POLLHUP | select.POLLERR
        else:
            mask = select.POLLHUP | select.POLLERR
        poll.register(in_fd, mask)

    res = poll.poll()

    for fd, signal in res:
        if signal & (select.POLLHUP | select.POLLERR):
            if fd == in_fd:
                # read the rest of the input.
                buf.append(input_fd.read())
                in_open = False
            elif fd == out_fd:
                out_open = False
        else:
            if fd == in_fd:
                data = input_fd.read()
                # recieved EOF, do not keep reading.
                if len(data) == 0:
                    in_open = False
                buf.append(data)
            elif fd == out_fd:
                try:
                    chars = output_fd.write(buf[0])
                except io.BlockingIOError, e:
                    # Try again with the characters that have not been written yet.
                    buf[0] = buf[0][e.characters_written:]
                else:
                    buf.pop(0)
        break