#!/usr/bin/python
# Copyright (C) 2005 Laurent A.V. Szyster
# http://svn.berlios.de/svnroot/repos/allegra/lib/netstring.py
#
# Modified by K. Vogel to use more sensible command names - verify instead
# of decode, decode instead of outline.  Also, output goes to stdout.
#
# Like many modules in Allegra, this one is its own first application and
# also provides a command line interface for shell scripts.
#
# Validate a stream of netstrings and assert the print of their count
# to STDERR:
#   ns verify < netstrings.net
#
# Decode a stream of netstrings and assert their count:
#   ns decode < netstrings.net 1> netlines.txt
#
# Encode a stream of lines as safe netstrings.  The CLI handles two more
# optional arguments for the verify and decode commands:
#   ns encode < lines.txt 1> safe.net
#
# The second optional argument is the size of the file buffer used, the
# last one is the maximum size of the decoder's buffer:
#   ns verify|decode [buffer] [max]
#
# Note that if not specified, the default command is decode, allowing
# the concise form:
#   ns < netstrings.net | more
#
# for its most common use.

class NetstringsError (Exception): pass

def encode (strings):
    "encode an sequence of 8-bit byte strings as netstrings"
    return ''.join (['%d:%s,' % (len (s), s) for s in strings])

def decode (buffer):
    "decode the netstrings found in the buffer, trunk garbage"
    size = len (buffer)
    prev = 0
    while prev < size:
        pos = buffer.find (':', prev)
        if pos < 1:
            break

        try:
            next = pos + int (buffer[prev:pos]) + 1
        except:
            break

        if next >= size:
            break

        if buffer[next] == ',':
            yield buffer[pos+1:next]

        else:
            break

        prev = next + 1

def validate (buffer, length):
    "decode the netstrings, but keep garbage and fit to size"
    size = len (buffer)
    prev = 0
    while prev < size and length:
        pos = buffer.find (':', prev)
        if pos < 1:
            if prev == 0:
                raise StopIteration # not a netstring!

            break

        try:
            next = pos + int (buffer[prev:pos]) + 1
        except:
            break

        if next >= size:
            break

        if buffer[next] == ',':
            length -= 1
            yield buffer[pos+1:next]

        else:
            break

        prev = next + 1

    if length:
        length -= 1
        yield buffer[max (prev, pos+1):]

        while length:
            length -= 1
            yield ''

def outline (encoded, format, indent):
    "recursively format nested netstrings as a CRLF outline"
    n = tuple (decode (encoded))
    if len (n) > 0:
        return ''.join ((outline (
            e, indent + format, indent
            ) for e in n))

    return format % encoded

def netstrings (instance):
    "encode a tree of instances as nested netstrings"
    t = type (instance)
    if t == str:
        return instance

    if t in (tuple, list, set, frozenset):
        return encode ((netstrings (i) for i in instance))

    if t == dict:
        return encode ((netstrings (i) for i in instance.items ()))

    try:
        return '%s' % instance

    except:
        return '%r' % instance

def netlist (encoded):
    "return a list of strings or [encoded] if no netstrings found"
    return list (decode (encoded)) or [encoded]

def nettree (encoded):
    "decode the nested encoded strings in a tree of lists"
    leaves = [nettree (s) for s in decode (encoded)]
    if len (leaves) > 0:
        return leaves

    return encoded

def netlines (encoded, format='%s\n', indent='  '):
    "beautify a netstring as an outline ready to log"
    n = tuple (decode (encoded))
    if len (n) > 0:
        return format % ''.join ((outline (
            e, format, indent
            ) for e in n))

    return format % encoded

def netoutline (encoded, indent=''):
    "recursively format nested netstrings as an outline with length"
    n = tuple (decode (encoded))
    if len (n) > 0:
        return '%s%d:\n%s%s,\n' % (
            indent, len (encoded), ''.join ((netoutline (
                e, indent + '  '
                ) for e in n)), indent)

    return '%s%d:%s,\n' % (indent, len (encoded), encoded)

def netpipe (more, BUFFER_MAX=0):
    """A practical netstrings pipe generator

    Decode the stream of netstrings produced by more (), raise
    a NetstringsError exception on protocol failure or StopIteration
    when the producer is exhausted.

    If specified, the BUFFER_MAX size must be more than twice as
    big as the largest netstring piped through (although netstrings
    strictly smaller than BUFFER_MAX may pass through without raising
    an exception).
    """
    buffer = more ()
    while buffer:
        pos = buffer.find (':')
        if pos < 0:
            raise NetstringsError, '1 not a netstring'
        try:
            next = pos + int (buffer[:pos]) + 1
        except:
            raise NetstringsError, '2 not a valid length'

        if 0 < BUFFER_MAX < next:
            raise (
                NetstringsError,
                '4 buffer overflow (%d bytes)' % BUFFER_MAX
                )

        while next >= len (buffer):
            data = more ()
            if data:
                buffer += data
            else:
                raise NetstringsError, '5 end of pipe'

        if buffer[next] == ',':
            yield buffer[pos+1:next]

        else:
            raise NetstringsError, '3 missing comma'

        buffer = buffer[next+1:]
        if buffer == '' or buffer.isdigit ():
            buffer += more ()

    # Note also that the first call to more must return at least the
    # encoded length of the first netstring, which practically is (or
    # should be) allways the case (for instance, piping in a netstring
    # sequence from a file will be done by blocks of pages, typically
    # between 512 and 4096 bytes, maybe more certainly not less).

if __name__ == '__main__':
    import sys

    if len (sys.argv) > 1:
        command = sys.argv[1]
    else:
        command = 'decode'

    if command in ('decode', 'verify'):
        if len (sys.argv) > 2:
            if len (sys.argv) > 3:
                try:
                    buffer_max = int (sys.argv[3])
                except:
                    sys.stderr.write ('3 invalid buffer max\n')
                    sys.exit (3)

            else:
                buffer_max = 0
            try:
                buffer_more = int (sys.argv[2])
            except:
                sys.stderr.write ('2 invalid buffer size\n')
                sys.exit (2)

        else:
            buffer_max = 0
            buffer_more = 4096
        def more ():
            return sys.stdin.read (buffer_more)
        count = 0
        try:
            if command == 'decode':
                write = sys.stdout.write
                for n in netpipe (more, buffer_max):
                    write (netlines (n))
                    count += 1
            else:
                for n in netpipe (more, buffer_max):
                    count += 1
        finally:
            assert None == sys.stdout.write ('%d\n' % count)
    elif command == 'encode':
        write = sys.stdout.write
        for line in sys.stdin.xreadlines ():
            write ('%d:%s,' % (len (line)-1, line[:-1]))
    else:
        sys.stderr.write ('1 invalid command\n')
        sys.exit (1)

    sys.exit (0)

# This library is free software; you can redistribute it and/or modify
# it under the terms of version 2 of the GNU General Public License as
# published by the Free Software Foundation.
#
#    http://www.gnu.org/copyleft/gpl.html
#
# This library is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#
# You should have received a copy of the GNU General Public License
# along with this library; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
# USA