#!/usr/bin/python3
#
# Simple LPD Server
#
# Based on RFC1179, this program will act as a simple LPD server, saving
# the data and control files to a preset directory (defaults to /tmp).
#
# Only the "Receive a printer job" command is supported.
#
########################################################################
#
# Changelog:
#
# 2020-12-04, bmason: fixed all the byte/str crap added by python 3
# 2016-05-12, bmason: Created
#
########################################################################

import socket
import os
import sys
import struct
import argparse
from time import sleep

lpd_commands = [
    'NULL',
    'Print any waiting jobs',
    'Receive a printer job',
    'Send queue state (short)',
    'Send queue state (long)',
    'Remove jobs']

receive_job_commands = [
    'Abort job',
    'Receive control file',
    'Receive data file']

job_errors = [
    'none',
    'down',
    'full',
    'bad']

job_error_descriptions = [
    'Success',
    'Queue not accepting jobs',
    'Queue temporarily full, retry later',
    'Bad job format, do not retry']

DEBUG = False

#------------------------------------------------------------------------------
class lpd():
    def __init__(self, socket, bufsiz):
        self.buf = b""
        self.bufsiz = bufsiz
        self.conn = socket
    
    def recv_command(self):
        if DEBUG: print('lpd.recv_command()', flush=True)
        nl = self.buf.find(b'\n')
        while nl == -1:
            addbuf = self.conn.recv(self.bufsiz)
            if len(addbuf) == 0:
                if DEBUG: print('    recv_command: no data.', flush=True)
                if DEBUG: print('lpd.recv_command() returning', repr(self.buf), flush=True)
                return self.buf
            
            self.buf = b''.join([self.buf, addbuf])
            nl = self.buf.find(b'\n')
            if DEBUG: print('    recv_command: buf=' + repr(self.buf), flush=True)
        rv = self.buf[:nl]
        self.buf = self.buf[nl+1:]
        if DEBUG: print('    recv_command: buf=' + repr(self.buf), flush=True)
        if DEBUG: print('lpd.recv_command() returning', repr(rv), flush=True)
        #return rv.encode('utf-8')
        return rv

    def recv_count(self, count):
        if DEBUG: print('lpd.recv_count(' + str(count) + ')', flush=True)
        if DEBUG: print('    recv_count: len(self.buf)=' + str(len(self.buf)), flush=True)
        if DEBUG: print('    recv_count: buf=' + repr(self.buf), flush=True)
        while len(self.buf) < count:
            addbuf = self.conn.recv(self.bufsiz)
            if len(addbuf) == 0:
                return self.buf
            
            self.buf = b''.join([self.buf, addbuf])
            if DEBUG: print('    recv_count: len(self.buf)='+str(len(self.buf)), flush=True)
            if DEBUG: print('    recv_count: buf=' + repr(self.buf), flush=True)

        rv = self.buf[:count]
        self.buf = self.buf[count:]
        if DEBUG: print('    recv_count: buf=' + repr(self.buf), flush=True)
        if DEBUG: print('lpd.recv_count() returning', repr(rv), flush=True)
        return rv

    def send(self, str):
        return self.conn.send(str)

    def shutdown(self, how):
        return self.conn.shutdown(how)

    def close(self):
        return self.conn.close()

#------------------------------------------------------------------------------
def receive_job(c):
    abort = False
    done = False
    file_type = ''

    while ( not done and not abort ):
        # Receive the subcommand (control or data file)
        try:
            subcmdstr = c.recv_command().decode('utf-8')
        except socket.error as msg:
            print('recv_command():', msg[1], flush=True)
            sys.exit(1)

        if DEBUG: print('subcommand:', repr(subcmdstr), flush=True)

        # If not data was received, then we're finished.
        if len(subcmdstr) == 0:
            done = True

        else:
            subcmd = subcmdstr[0]
            # Control file
            if subcmd == '\x02':
                size_str, file_name = subcmdstr[1:].split();
                file_size = int(size_str)
                file_type = 'Control'
                print('Receiving control file...', flush=True)
            # Data file
            elif subcmd == '\x03':
                size_str, file_name = subcmdstr[1:].split();
                file_size = int(size_str)
                file_type = 'Data'
                print('Receiving data file...', flush=True)
            # Abort
            elif subcmd == '\x01':
                abort = True
                file_name = ''
                file_type = ''
                file_size = 0
                print('Aborting job...',flush=True)
            # NULL for no reason
            elif subcmd == '\x00':
                file_name = ''
                file_type = 'Null'
                file_size = 0
                print('NULL...', flush=True)
            # Error!
            else:
                file_name = ''
                file_type = 'Err'
                file_size = 0
                print('Error...', flush=True)

            # Send response
            try:
                if file_type == 'Err':
                    print('NACK!', flush=True)
                    c.send(struct.pack('b', 4))
                else:
                    print('Ack.', flush=True)
                    c.send(struct.pack('b', 0))
                    
            except socket.error as msg:
                print('send():', msg[1])
                sys.exit(1)

            # Receive the {control,data} file, if any
            if file_size != 0:
                # Determine the width of the file_size field and
                # the string to print for updates.
                fsw = str(len(str(file_size)))
                if fsw == '': fsw = '1'
                update_str='{0} file {1}: {2:>'+fsw+'}/{3:>'+fsw+'} bytes\r'

                # Open file
                file_path = args.dest + '/' + file_name
                try:
                    f = open(file_path, 'w')
                except IOError as msg:
                    print('open():', msg[1])
                    sys.exit(1)
                    # Receive the file and save it

                # Read the entire file
                received = 0
                while received < file_size+1:
                    # The number of bytes to receive in this pass will be
                    # the smaller of bufsiz or the number of bytes
                    # remaining to be read.
                    remaining = file_size+1 - received
                    if args.bufsiz > remaining:
                        read_count = remaining
                    else:
                        read_count = args.bufsiz
                        
                    # Read from client
                    try:
                        buf = c.recv_count(read_count)
                    except socket.error as msg:
                        print('recv():', msg[1])
                        sys.exit(1)

                        if DEBUG: print('control/data:', repr(buf), flush=True)

                    # Write to file
                    try:
                        f.write(buf.decode())
                    except IOError as msg:
                        print('write():', msg[1])
                        sys.exit(1)
                    
                    received += len(buf)
                    print(update_str.format(
                        file_type, file_path, received, file_size), flush=True)

                    if args.wait > 0:
                        sleep(args.wait)
                # end while received < file_size+1
            
                # Close the file
                f.close()
                file_type = ''
                file_size = 0

                # Acknowledge file transfer
                print("Ack.", flush=True)
                try:
                    c.send(struct.pack('b', 0))
                except socket.error as msg:
                    print('send():', msg[1])
                    sys.exit(1)
                    
            # end if file_size != 0
        # end if len() == 0
    # end while (not done and not abort ):
# end receive_job()

#------------------------------------------------------------------------------
# Main program

# Don't buffer terminal outout
#sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 0)

# Process command line arguments
parser = argparse.ArgumentParser(description="Receive a print fie via LPD.")
parser.add_argument('-p', '--port', action='store', type=int,
                    default=515,
                    help='Port on which to listen (default is 515).')
parser.add_argument('-1', '--one', action='store_true',
                    help='Close the connection ane exit after one job.')
parser.add_argument('-b', '--bufsiz', action='store', type=int,
                    default=1024,
                    help='Size of the buffer to use (default is 4096).')
parser.add_argument('-d', '--dest', action='store', type=str,
                    default='/tmp',
                    help='Directory to store files in (default is "/tmp").')
parser.add_argument('-w', '--wait', action='store', type=int,
                    default=0,
                    help='Number of seconds to wait between calls '
                    'while receiving data.')
parser.add_argument('--joberror', action='store', type=str,
                    default='none', choices=job_errors,
                    help='Error to return in response to Receive Job. '
                    'Can be one of: ' 
                    + job_errors[1]+' (1 - '+job_error_descriptions[1]+'), ' 
                    + job_errors[2]+' (1 - '+job_error_descriptions[2]+'), ' 
                    + job_errors[3]+' (1 - '+job_error_descriptions[3]+')')
args = parser.parse_args()

# Make sure we're running as root.
if os.getuid() != 0 and args.port < 1025:
    print("Must be run as root to open privileged port " + str(args.port) + ".")
    sys.exit(1)
    

job_error_num = job_errors.index(args.joberror)

# Create a new socket
try:
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
except socket.error as msg:
    print('socket():',  msg[1])
    sys.exit(1)

# Allow socket re-use
try:
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
except socket.error as msg:
    print('setsockopt():',  msg[1])
    sys.exit(1)
   
# Bind the socket to the port
try:
    s.bind(('', args.port))
except socket.error as msg:
    print('bind(\'\',', args.port,  '):',  msg[1])
    sys.exit(1)

# Do everything.  Catch ^C and exit cleanly(ish)
try:
    done = False
    while not done:
        # Listen for an incoming connection
        print('Listening on ' + str(args.port) + '...', flush=True)
        try:
            s.listen(1)
        except socket.error as msg:
            print('listen():', msg[1])
            sys.exit(1)

        # Got a connection, accept it and create an lpd connection with it.
        conn, addr = s.accept()
        c = lpd(conn, args.bufsiz)
        print('New connection from', str(addr[0]) + ':' + str(addr[1]), flush=True)

        # Continue to accept commands as long as we're connected
        is_connected = True
        while is_connected:
            try:
                cmdstr = c.recv_command()
            except socket.error as msg:
                print('recv_command():', msg[1])
                sys.exit(1)

            # If we didn't receive any data, then we're disconnected
            if len(cmdstr) == 0:
                is_connected = False
            else:
                if job_error_num == 0:
                    # Separate the command (first octet) from the rest of the
                    # command
                    cmd,rest = struct.unpack('b'+str(len(cmdstr)-1)+'s', cmdstr)

                    # Receive a printer job
                    if cmd == 2:
                        queue = cmdstr[1:]
                        if queue == '': queue = '(empty)'
                        print('Receiving job for', queue.decode('utf-8'), '...', flush=True)
                    
                        # Send acknowledgement
                        print('Ack.', flush=True)
                        try:
                            c.send(struct.pack('b', 0))
                        except socket.error as msg:
                            print('send():', msg[1])
                            sys.exit(1)
            
                        # Receive the print job
                        receive_job(c)
                        
                    # Everything else
                    else:
                        # All other functions are invalid, so send error.
                        print('Function ' + str(cmd) , flush=True)
                        if cmd <= len(lpd_commands):
                            print(' (' + lpd_commands[cmd] + ')' , flush=True)
                            print(' not implemented.', flush=True)
                            try:
                                c.send(struct.pack('b',3))
                            except socket.error as msg:
                                print('send():', msg[1])
                                sys.exit(1)
                    # end if cmd==2/else
                    
                # Return an error code
                else:
                    print('Sending error' , str(job_error_num) \
                    + ' ("' + job_error_descriptions[job_error_num] \
                    + '")', flush=True)
                    try:
                        c.send(struct.pack('b',job_error_num))
                    except socket.error as msg:
                        print('send():', msg[1])
                        sys.exit(1)
                        
                # end if args.joberror == ''/else
            # end if len(cmdstr) == 0
        # end while is_connected

        # Close the connection
        print('Closing Connection from', str(addr[0]) + ':' + str(addr[1]), flush=True)
        try:
            c.shutdown(socket.SHUT_RDWR)
            c.close()
        except NameError:
            print
        except socket.error:
            print

        # If we're only doing this once, then were done after one loop.
        if args.one:
            done = True
    # end while not_done

# Handle keyboard interrrupt.
except KeyboardInterrupt:
    print('Keyboard Interrupt', flush=True)

print('Closing listening socket...', flush=True)

try:
    s.shutdown(socket.SHUT_RDWR)
    s.close()
except NameError:
    print

# END

