upm

#!/usr/local/bin/python3

import base64
import gzip
import io
import json
import os
import queue
import re
import signal
import struct
import subprocess
import sys
import threading
import traceback

import requests

VERSION = "1.1.2019.509"
SCHEMA = "upm"

PVG_MAGIC = b"PVG\x00\x00\x00\x00\x00\x01\x00\x00\x00"
PNG_MAGIC = b"\x89PNG\x0d\x0a\x1a\x0a"

CUPS_BACKEND_OK = 0  # Job completed successfully
CUPS_BACKEND_FAILED = 1  # Job failed, use error-policy
CUPS_BACKEND_AUTH_REQUIRED = 2  # Job failed, authentication required
CUPS_BACKEND_HOLD = 3  # Job failed, hold job
CUPS_BACKEND_STOP = 4  # Job failed, stop queue
CUPS_BACKEND_CANCEL = 5  # Job failed, cancel job
CUPS_BACKEND_RETRY = 6  # Job failed, retry this job later
CUPS_BACKEND_RETRY_CURRENT = 7  # Job failed, retry this job immediately

BUF_SIZE = 1024 * 64  # 64K read buffer
QUEUE_LENGTH = 1024  # 1024 * BUF_SIZE for compression

os.environ["PATH"] = "/usr/local/bin" + os.pathsep + os.environ["PATH"]
args = len(sys.argv)


class ServerException(Exception):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class RequestException(Exception):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class BackendStopException(Exception):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


def sigterm(signum, frame):
    raise BackendStopException()


def log(msg, level):
    print(f"{level}: {msg}", file=sys.stderr)


class Popen(threading.Thread):
    def __init__(self, args, stdin):
        super().__init__()
        self.proc = subprocess.Popen(args, stdin=stdin, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
        self.stdin = stdin
        self.data = None
        self.start()

    def run(self):
        self.data = self.proc.stdout.read()

    def join(self, timeout):
        super().join(timeout)
        if self.is_alive():
            self.proc.terminate()
            super().join()
        else:
            self.proc.wait(timeout)


class Compress(threading.Thread):
    def __init__(self, fileobj):
        super().__init__()
        self.gf = gzip.GzipFile(mode='wb', fileobj=fileobj)
        self.buf_queue = queue.Queue(QUEUE_LENGTH)

        self.start()

    def run(self):
        while True:
            buf = self.buf_queue.get()
            if not buf:
                break
            self.gf.write(buf)

    def close(self):
        self.buf_queue.put(False)
        self.join()
        self.gf.close()
        del self.buf_queue

    def write(self, data):
        self.buf_queue.put(data)


class APIRequest:
    def __init__(self, base_url):
        self.s = requests.session()
        self.base = base_url
        self.s.headers["Referer"] = base_url + "/driver/print.html"
        self.s.headers["User-Agent"] = "Mozilla/4.0 (compatible; MSIE 7.0; Windows NT 6.2; WOW64; Trident/7.0; .NET4.0C; .NET4.0E; .NET CLR 2.0.50727; .NET CLR 3.0.30729; .NET CLR 3.5.30729; LCTE)"

    def __parse_result(self, resp):
        if not resp.status_code == 200:
            raise ServerException(resp.request.url)

        content = json.loads(resp.content.decode())
        if not content["code"] == 0:
            raise RequestException(content["message"])

        return content

    def request(self, suburl, data):
        with self.s.post(self.base + suburl, json=data, timeout=10) as resp:
            return self.__parse_result(resp)

    def upload(self, suburl, stream, filename):
        with self.s.post(self.base + suburl, files={"szFile": (filename, stream, "application/octet-stream", {"Content-Transfer-Encoding": "binary"})}, timeout=180) as resp:
            return self.__parse_result(resp)

    def getSID(self):
        if "OSESSIONID" in self.s.cookies:
            return self.s.cookies["OSESSIONID"]
        else:
            self.request("/api/client/Paper/GetPaper", data={})
            return self.getSID()

    def close(self):
        self.s.close()


def start_print(stream):
    device_uri = os.environ["DEVICE_URI"]
    job_title = sys.argv[3]

    parse = re.match(fr"{SCHEMA}://(.*)/(.*?):(.*)", device_uri)
    if parse is None:
        log("Invalid device URI", "ERROR")
        sys.exit(CUPS_BACKEND_STOP)

    username = parse.groups()[1]
    password = parse.groups()[2]
    url = "http://" + parse.groups()[0]
    log(f"Username: {username}", "INFO")
    log(f"Password: {password}", "DEBUG")
    log(f"URL: {url}", "INFO")

    client = APIRequest(url)
    jobID = None
    try:
        signal.signal(signal.SIGTERM, sigterm)

        sid = client.getSID()

        log("Start Login", "INFO")
        auth = client.request("/api/client/Auth/GetAuthToken", data={"sessionid": sid})
        login = client.request("/api/client/Auth/Login", data={"szLogonName": username, "szPassword": password, "szToken": auth["szToken"], "sessionid": sid})
        client.request("/api/client/Auth/Cancel", data={"sessionid": sid, "szToken": auth["szToken"]})
        log(f"Print for {login['result']['szTrueName']}", "INFO")

        write = []
        gdata = io.BytesIO()
        gf = Compress(fileobj=gdata)
        write.append(gf)

        pvg_pr, pvg_pw = os.pipe()
        cl_pr, cl_pw = os.pipe()
        pvg_proc = None
        cl_proc = None
        try:
            pvg_proc = Popen(["gpcl6", "-dSAFER", "-dNOPAUSE", "-sDEVICE=pngalpha", "-sOutputFile=%stdout", "-"], stdin=pvg_pr)
            write.append(open(pvg_pw, 'wb'))
            cl_proc = Popen(["gpcl6", "-dSAFER", "-dNOPAUSE", "-sDEVICE=inkcov", "-sOutputFile=%stdout", "-"], stdin=cl_pr)
            write.append(open(cl_pw, 'wb'))
        except OSError as e:
            log(f"Failed to get printing info: {str(e)}", "WARNING")
        finally:
            os.close(pvg_pr)
            os.close(cl_pr)

        while True:
            buf = stream.read(BUF_SIZE)
            if buf == b'':
                break
            try:
                for w in write:
                    w.write(buf)
            except BrokenPipeError:
                log("Failed to write to pipe", "WARNING")
                w.close()
                write.remove(w)
        for w in write:
            w.close()

        log("Generate Page Info", "INFO")
        preview = PVG_MAGIC
        if pvg_proc:
            pvg_proc.join(300)
            png_data = pvg_proc.data

            ind = 0
            pages = 0
            suffix = b''
            png_data = png_data + b"\x00"
            while True:
                ind_new = png_data.find(PNG_MAGIC, ind + 1)

                png = base64.b64encode(png_data[ind:ind_new])
                suffix = suffix + struct.pack("<I", len(preview))
                preview = preview + struct.pack("<I", len(png)) + png

                pages = pages + 1
                if ind_new == -1:
                    break
                ind = ind_new
            preview = preview + suffix + struct.pack("<I", pages)

        if cl_proc:
            cl_proc.join(120)
            cmyk = cl_proc.data.decode().split("\n")

            color_map = []
            for line in cmyk:
                if not line:
                    continue

                color = re.match(r" ([\d\.]+)  ([\d\.]+)  ([\d\.]+)  ([\d\.]+) CMYK OK", line)
                if color is None:
                    log(f"Invalid output from 'inkcov': {line}", "WARNING")
                    continue
                c = float(color.groups()[0])
                m = float(color.groups()[1])
                y = float(color.groups()[2])
                k = float(color.groups()[3])

                if c == m == y == k == 0:
                    continue
                elif c == m == y:
                    color_map.append(0)
                else:
                    color_map.append(1)
            v_pages = len(color_map)

            if pages - v_pages > 0:
                log(f"Skipped {pages - v_pages} empty page(s)", "NOTICE")

            if v_pages == 0:
                log("No valid pages found! Set to one", "WARNING")
                color_map = [0]
                v_pages = len(color_map)

            bw_pages = color_map.count(0)
            color_pages = v_pages - bw_pages
        else:
            color_map = [0]
            bw_pages = 1
            color_pages = 0
            v_pages = 1

        log(f"Create Job [Valid Pages: {v_pages}, Color: {color_pages}, Grayscale: {bw_pages}]", "INFO")
        # TODO: Add A3 paper and single/duplex supports. Single is default now
        # TODO: Add "切换为黑白输出" support
        job_spec = {
            "dwProperty": 0,
            "szJobName": job_title,
            "dwCopies": 1,
            "szAttribe": ("" if color_pages == 0 else "color,") + "single,collate,NUP1,",
            "szPaperDetail": json.dumps([{"dwPaperID": 9, "dwBWPages": bw_pages, "dwColorPages": color_pages, "dwPaperNum": v_pages}]),
            "szColorMap": "".join(str(i) for i in color_map),
            "OSESSIONID": sid
        }
        job = client.request("/api/client/PrintJob/Create", data=job_spec)
        jobID = job['result']['dwJobId']

        log("Upload Job", "INFO")
        gdata = gdata.getvalue()
        client.upload(f"/api/client/PrintJob/Upload?dwJobId={jobID}", gdata, "raw.dat")
        client.upload(f"/api/client/PrintJob/UploadPreview?dwJobId={jobID}", preview, "preview.pvg")

        client.request("/api/client/PrintJob/Set", data={"dwJobId": jobID, "dwStatus": 1, "OSESSIONID": sid})
        log("Printing Completed", "INFO")
        jobID = None
    except ServerException as e:
        log(f"Printing Server Error in [{str(e)}]", "ERROR")
        sys.exit(CUPS_BACKEND_RETRY)

    except RequestException as e:
        log(f"Invalid Request: {str(e)}", "ERROR")
        sys.exit(CUPS_BACKEND_HOLD)

    except requests.RequestException:
        log(f"Connect to Server Failed", "ERROR")
        sys.exit(CUPS_BACKEND_RETRY)

    except BackendStopException:
        log(f"Printing Canceled", "INFO")
        sys.exit(CUPS_BACKEND_FAILED)

    except Exception as e:
        log(f"Backend Error: {str(e)}", "ERROR")
        tb = traceback.format_exc()
        [log(f"{i}", "ERROR") for i in tb.split("\n")]
        sys.exit(CUPS_BACKEND_STOP)

    finally:
        if jobID is not None:
            try:
                client.request("/api/client/PrintJob/Del", data={"dwJobId": jobID})
            except Exception as e:
                log(f"Failed to Remove Printing Job: {str(e)}", "ERROR")
        client.close()


if __name__ == "__main__":
    if args == 1:
        print(f"network {SCHEMA} \"Unifound\" \"Unifound Printing Backend [{VERSION}]\"")
        sys.exit(CUPS_BACKEND_OK)
    elif args == 6:
        start_print(sys.stdin.buffer)
    elif args == 7:
        with open(sys.argv[6], 'rb') as spool_file:
            start_print(spool_file)
    else:
        print("Not implemented")
        sys.exit(CUPS_BACKEND_FAILED)