#!/usr/bin/env python
# This file is part of Xpra.
# Copyright (C) 2011-2015 Antoine Martin <antoine@devloop.org.uk>
# Copyright (C) 2008, 2009, 2010 Nathaniel Smith <njs@pobox.com>
# Xpra is released under the terms of the GNU GPL v2, or, at your option, any
# later version. See the file COPYING for details.

import sys
import zlib

from xpra.log import Logger
log = Logger("network", "protocol")
from xpra.net.header import LZ4_FLAG, ZLIB_FLAG, LZO_FLAG
from xpra.os_util import _memoryview


python_lz4_version = None
lz4_version = None
try:
    import lz4
    try:
        from lz4 import VERSION as lz4_VERSION                          #@UnresolvedImport
    except:
        #older versions (0.7 and older):
        import pkg_resources
        lz4_VERSION = pkg_resources.get_distribution("lz4").version
    from lz4 import LZ4_compress, compressHC, LZ4_uncompress        #@UnresolvedImport
    has_lz4 = True
    if hasattr(lz4, "LZ4_compress_fast"):
        from lz4 import LZ4_compress_fast                           #@UnresolvedImport
        def lz4_compress(packet, level):
            if level>=9:
                return level | LZ4_FLAG, compressHC(packet)
            #clamp it: 0->17, 1->12, 2->7, 3->2, >=4->1
            accel = max(1, 17-level*5)
            return level | LZ4_FLAG, LZ4_compress_fast(packet, accel)
    else:
        #v0.7.0 and earlier
        def lz4_compress(packet, level):
            if level>=9:
                return level | LZ4_FLAG, compressHC(packet)
            return level | LZ4_FLAG, LZ4_compress(packet)
    #try to figure out the version number:
    python_lz4_version = lz4_VERSION
    assert python_lz4_version>="0.7", "python-lz4 version %s is older than 0.7.0 are vulnerable and should not be used, see CVE-2014-4715" % lz4_version
    #now try to check the underlying "liblz4" version
    #which is only available with python-lz4 0.8.0 onwards:
    if hasattr(lz4, "LZ4_VERSION"):
        try:
            from distutils.version import LooseVersion
        except ImportError:
            pass
        else:
            #last known security issue:
            assert LooseVersion(lz4.LZ4_VERSION)>=LooseVersion("r119"), "lz4 version %s is vulnerable and should not be used, see CVE-2014-4715" % lz4.LZ4_VERSION
        lz4_version = lz4.LZ4_VERSION
except Exception as e:
    log("lz4 not found: %s", e)
    LZ4_uncompress = None
    has_lz4 = False
    def lz4_compress(packet, level):
        raise Exception("lz4 is not supported!")

python_lzo_version = None
lzo_version = None
try:
    import lzo
    has_lzo = True
    python_lzo_version = lzo.__version__
    lzo_version = lzo.LZO_VERSION_STRING
    def lzo_compress(packet, level):
        return level | LZO_FLAG, lzo.compress(packet)
    LZO_decompress = lzo.decompress
except Exception as e:
    log("lzo not found: %s", e)
    LZO_decompress = None
    has_lzo = False
    def lzo_compress(packet, level):
        raise Exception("lzo is not supported!")


#stupid python version breakage:
if sys.version > '3':
    def zcompress(packet, level):
        if type(packet)!=bytes:
            packet = bytes(packet, 'UTF-8')
        return level + ZLIB_FLAG, zlib.compress(packet, level)

    def nocompress(packet, level):
        if type(packet)!=bytes:
            packet = bytes(packet, 'UTF-8')
        return 0, packet
else:
    def zcompress(packet, level):
        return level + ZLIB_FLAG, zlib.compress(str(packet), level)
    def nocompress(packet, level):
        return 0, packet

#defaults to True if available:
use_zlib = True
use_lzo = has_lzo
use_lz4 = has_lz4

#all the compressors we know about, in best compatibility order:
ALL_COMPRESSORS = ["zlib", "lz4", "lzo"]

#order for performance:
PERFORMANCE_ORDER = ["lz4", "lzo", "zlib"]


_COMPRESSORS = {
        "zlib"  : zcompress,
        "lz4"   : lz4_compress,
        "lzo"   : lzo_compress,
        "none"  : nocompress,
               }

def get_compression_caps():
    caps = {
            "lz4"                   : use_lz4,
            "lzo"                   : use_lzo,
            "zlib"                  : use_zlib,
            "zlib.version"          : zlib.__version__,
           }
    if lzo_version:
        caps["lzo.version"] = lzo_version
    if python_lzo_version:
        caps["python-lzo.version"] = python_lzo_version
    if lz4_version:
        caps["lz4.version"] = lz4_version
    if python_lz4_version:
        caps["python-lz4.version"] = python_lz4_version
    return caps

def get_enabled_compressors(order=ALL_COMPRESSORS):
    enabled = [x for x,b in {
            "lz4"                   : use_lz4,
            "lzo"                   : use_lzo,
            "zlib"                  : use_zlib,
            }.items() if b]
    #order them:
    return [x for x in order if x in enabled]

def get_compressor(c):
    assert c=="none" or c in ALL_COMPRESSORS
    return _COMPRESSORS[c]

def get_compressor_name(c):
    assert c in _COMPRESSORS.values(), "invalid compressor: %s" % c
    for k,v in _COMPRESSORS.items():
        if v==c:
            return k
    raise Exception("impossible bug!")


def sanity_checks():
    if not use_lzo and not use_lz4:
        if not use_zlib:
            log.warn("Warning: all the compressors are disabled,")
            log.warn(" unless you have a gigabit connection or better, performance will suffer")
        else:
            log.warn("Warning: zlib is the only compressor enabled")
            log.warn(" install and enable lzo or lz4 support for better performance")


class Compressed(object):
    def __init__(self, datatype, data, can_inline=False):
        self.datatype = datatype
        self.data = data
        self.can_inline = can_inline
    def __len__(self):
        return len(self.data)
    def __repr__(self):
        return  "Compressed(%s: %i bytes)" % (self.datatype, len(self.data))


class LevelCompressed(Compressed):
    def __init__(self, datatype, data, level, algo, can_inline):
        Compressed.__init__(self, datatype, data, can_inline)
        self.level = level
        self.algorithm = algo
    def __repr__(self):
        return  "LevelCompressed(%s: %i bytes as %s/%i)" % (self.datatype, len(self.data), self.algorithm, self.level)


class Uncompressed(object):
    def __init__(self, datatype, data):
        self.datatype = datatype
        self.data = data
    def __len__(self):
        return len(self.data)
    def __repr__(self):
        return  "Uncompressed(%s: %i bytes)" % (self.datatype, len(self.data))
    def compress(self):
        raise Exception("compress() not defined on %s" % self)

def compressed_wrapper(datatype, data, level=5, zlib=False, lz4=False, lzo=False, can_inline=True):
    if _memoryview and isinstance(data, _memoryview):
        data = data.tobytes()
    if lz4:
        assert use_lz4, "cannot use lz4"
        algo = "lz4"
        cl, cdata = lz4_compress(data, level)
    elif lzo:
        assert use_lzo, "cannot use lzo"
        algo = "lzo"
        cl, cdata = lzo_compress(data, level)
    else:
        assert use_zlib, "cannot use zlib"
        algo = "zlib"
        cl, cdata = zcompress(data, level)
    return LevelCompressed(datatype, cdata, cl, algo, can_inline=can_inline)


class InvalidCompressionException(Exception):
    pass


def get_compression_type(level):
    if level & LZ4_FLAG:
        return "lz4"
    elif level & LZO_FLAG:
        return "lzo"
    else:
        return "zlib"

import struct
LZ4_HEADER = struct.Struct('<L')
def decompress(data, level):
    #log.info("decompress(%s bytes, %s) type=%s", len(data), get_compression_type(level))
    if level & LZ4_FLAG:
        if not has_lz4:
            raise InvalidCompressionException("lz4 is not available")
        if not use_lz4:
            raise InvalidCompressionException("lz4 is not enabled")
        size = LZ4_HEADER.unpack_from(data[:4])[0]
        #TODO: it would be better to use the max_size we have in protocol,
        #but this hardcoded value will do for now
        assert size<=(256*1024*1024)
        return LZ4_uncompress(data)
    elif level & LZO_FLAG:
        if not has_lzo:
            raise InvalidCompressionException("lzo is not available")
        if not use_lzo:
            raise InvalidCompressionException("lzo is not enabled")
        return LZO_decompress(data)
    else:
        if not use_zlib:
            raise InvalidCompressionException("zlib is not enabled")
        return zlib.decompress(data)

NAME_TO_FLAG = {
                "lz4"   : LZ4_FLAG,
                "zlib"  : 0,
                "lzo"   : LZO_FLAG,
                }

def decompress_by_name(data, algo):
    assert algo in NAME_TO_FLAG, "invalid compression algorithm: %s" % algo
    flag = NAME_TO_FLAG[algo]
    return decompress(data, flag)


def main():
    from xpra.platform import init, clean
    try:
        init("Compression", "Compression Info")
        for k,v in sorted(get_compression_caps().items()):
            print(k.ljust(20)+": "+str(v))
    finally:
        #this will wait for input on win32:
        clean()


if __name__ == "__main__":
    main()
