#
# Copyright (c) 2014 Sylvain Peyrefitte
#
# This file is part of rdpy.
#
# rdpy is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program 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. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#

"""
Definition of structure use for capabilities nego
Use in PDU layer
"""

from rdpy.network.type import CompositeType, String, UInt8, UInt16Le, UInt32Le, sizeof, ArrayType, FactoryType
    
class CapsType(object):
    """
    Different type of capabilities
    @see: http://msdn.microsoft.com/en-us/library/cc240486.aspx
    """
    CAPSTYPE_GENERAL = 0x0001
    CAPSTYPE_BITMAP = 0x0002
    CAPSTYPE_ORDER = 0x0003
    CAPSTYPE_BITMAPCACHE = 0x0004
    CAPSTYPE_CONTROL = 0x0005
    CAPSTYPE_ACTIVATION = 0x0007
    CAPSTYPE_POINTER = 0x0008
    CAPSTYPE_SHARE = 0x0009
    CAPSTYPE_COLORCACHE = 0x000A
    CAPSTYPE_SOUND = 0x000C
    CAPSTYPE_INPUT = 0x000D
    CAPSTYPE_FONT = 0x000E
    CAPSTYPE_BRUSH = 0x000F
    CAPSTYPE_GLYPHCACHE = 0x0010
    CAPSTYPE_OFFSCREENCACHE = 0x0011
    CAPSTYPE_BITMAPCACHE_HOSTSUPPORT = 0x0012
    CAPSTYPE_BITMAPCACHE_REV2 = 0x0013
    CAPSTYPE_VIRTUALCHANNEL = 0x0014
    CAPSTYPE_DRAWNINEGRIDCACHE = 0x0015
    CAPSTYPE_DRAWGDIPLUS = 0x0016
    CAPSTYPE_RAIL = 0x0017
    CAPSTYPE_WINDOW = 0x0018
    CAPSETTYPE_COMPDESK = 0x0019
    CAPSETTYPE_MULTIFRAGMENTUPDATE = 0x001A
    CAPSETTYPE_LARGE_POINTER = 0x001B
    CAPSETTYPE_SURFACE_COMMANDS = 0x001C
    CAPSETTYPE_BITMAP_CODECS = 0x001D
    CAPSSETTYPE_FRAME_ACKNOWLEDGE = 0x001E
    
class MajorType(object):
    """
    Use in general capability
    @see: http://msdn.microsoft.com/en-us/library/cc240549.aspx
    """
    OSMAJORTYPE_UNSPECIFIED = 0x0000
    OSMAJORTYPE_WINDOWS = 0x0001
    OSMAJORTYPE_OS2 = 0x0002
    OSMAJORTYPE_MACINTOSH = 0x0003
    OSMAJORTYPE_UNIX = 0x0004
    OSMAJORTYPE_IOS = 0x0005
    OSMAJORTYPE_OSX = 0x0006
    OSMAJORTYPE_ANDROID = 0x0007
        
class MinorType(object):
    """
    Use in general capability
    @see: http://msdn.microsoft.com/en-us/library/cc240549.aspx
    """
    OSMINORTYPE_UNSPECIFIED = 0x0000
    OSMINORTYPE_WINDOWS_31X = 0x0001
    OSMINORTYPE_WINDOWS_95 = 0x0002
    OSMINORTYPE_WINDOWS_NT = 0x0003
    OSMINORTYPE_OS2_V21 = 0x0004
    OSMINORTYPE_POWER_PC = 0x0005
    OSMINORTYPE_MACINTOSH = 0x0006
    OSMINORTYPE_NATIVE_XSERVER = 0x0007
    OSMINORTYPE_PSEUDO_XSERVER = 0x0008
    OSMINORTYPE_WINDOWS_RT = 0x0009
 
class GeneralExtraFlag(object):
    """
    Use in general capability
    @see: http://msdn.microsoft.com/en-us/library/cc240549.aspx
    """
    FASTPATH_OUTPUT_SUPPORTED = 0x0001
    NO_BITMAP_COMPRESSION_HDR = 0x0400
    LONG_CREDENTIALS_SUPPORTED = 0x0004
    AUTORECONNECT_SUPPORTED = 0x0008
    ENC_SALTED_CHECKSUM = 0x0010
      
class Boolean(object):
    FALSE = 0x00
    TRUE = 0x01
 
class OrderFlag(object):
    """
    Use in order capability
    @see: http://msdn.microsoft.com/en-us/library/cc240556.aspx
    """
    NEGOTIATEORDERSUPPORT = 0x0002
    ZEROBOUNDSDELTASSUPPORT = 0x0008
    COLORINDEXSUPPORT = 0x0020
    SOLIDPATTERNBRUSHONLY = 0x0040
    ORDERFLAGS_EXTRA_FLAGS = 0x0080
     
class Order(object):
    """
    Drawing orders supported
    Use in order capability
    @see: http://msdn.microsoft.com/en-us/library/cc240556.aspx
    """
    TS_NEG_DSTBLT_INDEX = 0x00
    TS_NEG_PATBLT_INDEX = 0x01
    TS_NEG_SCRBLT_INDEX = 0x02
    TS_NEG_MEMBLT_INDEX = 0x03
    TS_NEG_MEM3BLT_INDEX = 0x04
    TS_NEG_DRAWNINEGRID_INDEX = 0x07
    TS_NEG_LINETO_INDEX = 0x08
    TS_NEG_MULTI_DRAWNINEGRID_INDEX = 0x09
    TS_NEG_SAVEBITMAP_INDEX = 0x0B
    TS_NEG_MULTIDSTBLT_INDEX = 0x0F
    TS_NEG_MULTIPATBLT_INDEX = 0x10
    TS_NEG_MULTISCRBLT_INDEX = 0x11
    TS_NEG_MULTIOPAQUERECT_INDEX = 0x12
    TS_NEG_FAST_INDEX_INDEX = 0x13
    TS_NEG_POLYGON_SC_INDEX = 0x14
    TS_NEG_POLYGON_CB_INDEX = 0x15
    TS_NEG_POLYLINE_INDEX = 0x16
    TS_NEG_FAST_GLYPH_INDEX = 0x18
    TS_NEG_ELLIPSE_SC_INDEX = 0x19
    TS_NEG_ELLIPSE_CB_INDEX = 0x1A
    TS_NEG_INDEX_INDEX = 0x1B
        
class OrderEx(object):
    """
    Extension orders
    Use in order capability
    """
    ORDERFLAGS_EX_CACHE_BITMAP_REV3_SUPPORT = 0x0002
    ORDERFLAGS_EX_ALTSEC_FRAME_MARKER_SUPPORT = 0x0004

class InputFlags(object):
    """
    Input flag use in input capability
    @see:  http://msdn.microsoft.com/en-us/library/cc240563.aspx
    """
    INPUT_FLAG_SCANCODES = 0x0001
    INPUT_FLAG_MOUSEX = 0x0004
    INPUT_FLAG_FASTPATH_INPUT = 0x0008
    INPUT_FLAG_UNICODE = 0x0010
    INPUT_FLAG_FASTPATH_INPUT2 = 0x0020
    INPUT_FLAG_UNUSED1 = 0x0040
    INPUT_FLAG_UNUSED2 = 0x0080
    TS_INPUT_FLAG_MOUSE_HWHEEL = 0x0100

class BrushSupport(object):
    """
    Brush support of client
    @see: http://msdn.microsoft.com/en-us/library/cc240564.aspx
    """
    BRUSH_DEFAULT = 0x00000000
    BRUSH_COLOR_8x8 = 0x00000001
    BRUSH_COLOR_FULL = 0x00000002

class GlyphSupport(object):
    """
    Use by glyph order
    @see: http://msdn.microsoft.com/en-us/library/cc240565.aspx
    """
    GLYPH_SUPPORT_NONE = 0x0000
    GLYPH_SUPPORT_PARTIAL = 0x0001
    GLYPH_SUPPORT_FULL = 0x0002
    GLYPH_SUPPORT_ENCODE = 0x0003
   
class OffscreenSupportLevel(object):
    """
    Use to determine offscreen cache level supported
    @see: http://msdn.microsoft.com/en-us/library/cc240550.aspx
    """
    FALSE = 0x00000000
    TRUE = 0x00000001
 
class VirtualChannelCompressionFlag(object):
    """
    Use to determine virtual channel compression
    @see: http://msdn.microsoft.com/en-us/library/cc240551.aspx
    """
    VCCAPS_NO_COMPR = 0x00000000
    VCCAPS_COMPR_SC = 0x00000001
    VCCAPS_COMPR_CS_8K = 0x00000002
  
class SoundFlag(object):
    """
    Use in sound capability to inform it
    @see: http://msdn.microsoft.com/en-us/library/cc240552.aspx
    """
    NONE = 0x0000
    SOUND_BEEPS_FLAG = 0x0001

class CacheEntry(CompositeType):
    """
    Use in capability cache exchange
    @see: http://msdn.microsoft.com/en-us/library/cc240566.aspx
    """
    def __init__(self):
        CompositeType.__init__(self)
        self.cacheEntries = UInt16Le()
        self.cacheMaximumCellSize = UInt16Le()
    
    
class Capability(CompositeType):
    """
    A capability
    @see: http://msdn.microsoft.com/en-us/library/cc240486.aspx
    """
    def __init__(self, capabilitySetType = 0, capability = None):
        CompositeType.__init__(self)
        self.capabilitySetType = UInt16Le(capabilitySetType, constant = (not capability is None))
        self.lengthCapability = UInt16Le(lambda:sizeof(self))
        
        def CapabilityFactory():
            '''
            closure for capability factory
            '''
            if self.capabilitySetType.value == CapsType.CAPSTYPE_GENERAL:
                return GeneralCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_BITMAP:
                return BitmapCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_ORDER:
                return OrderCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_BITMAPCACHE:
                return BitmapCacheCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_POINTER:
                return PointerCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_INPUT:
                return InputCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_BRUSH:
                return BrushCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_GLYPHCACHE:
                return GlyphCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_OFFSCREENCACHE:
                return OffscreenBitmapCacheCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_VIRTUALCHANNEL:
                return VirtualChannelCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_SOUND:
                return SoundCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_CONTROL:
                return ControlCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_ACTIVATION:
                return WindowActivationCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_FONT:
                return FontCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_COLORCACHE:
                return ColorCacheCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            elif self.capabilitySetType.value == CapsType.CAPSTYPE_SHARE:
                return ShareCapability(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
            else:
                return String(readLen = UInt16Le(lambda:self.lengthCapability.value - 4))
        
        if capability is None:
            capability = CapabilityFactory
            
        self.capability = FactoryType(capability)

class GeneralCapability(CompositeType):
    """
    General capability (protocol version and compression mode)
    client -> server
    server -> client
    @see: http://msdn.microsoft.com/en-us/library/cc240549.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.osMajorType = UInt16Le()
        self.osMinorType = UInt16Le()
        self.protocolVersion = UInt16Le(0x0200, constant = True)
        self.pad2octetsA = UInt16Le()
        self.generalCompressionTypes = UInt16Le(0, constant = True)
        self.extraFlags = UInt16Le()
        self.updateCapabilityFlag = UInt16Le(0, constant = True)
        self.remoteUnshareFlag = UInt16Le(0, constant = True)
        self.generalCompressionLevel = UInt16Le(0, constant = True)
        self.refreshRectSupport = UInt8()
        self.suppressOutputSupport = UInt8()
        
class BitmapCapability(CompositeType):
    """
    Bitmap format Capability
    client -> server
    server -> client
    @see: http://msdn.microsoft.com/en-us/library/cc240554.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.preferredBitsPerPixel = UInt16Le()
        self.receive1BitPerPixel = UInt16Le(0x0001)
        self.receive4BitsPerPixel = UInt16Le(0x0001)
        self.receive8BitsPerPixel = UInt16Le(0x0001)
        self.desktopWidth = UInt16Le()
        self.desktopHeight = UInt16Le()
        self.pad2octets = UInt16Le()
        self.desktopResizeFlag = UInt16Le()
        self.bitmapCompressionFlag = UInt16Le(0x0001, constant = True)
        self.highColorFlags = UInt8(0)
        self.drawingFlags = UInt8()
        self.multipleRectangleSupport = UInt16Le(0x0001, constant = True)
        self.pad2octetsB = UInt16Le()
        
class OrderCapability(CompositeType):
    """
    Order capability list all drawing order supported
    client -> server
    server -> client
    @see: http://msdn.microsoft.com/en-us/library/cc240556.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.terminalDescriptor = String("\x00" * 16, readLen = UInt8(16))
        self.pad4octetsA = UInt32Le(0)
        self.desktopSaveXGranularity = UInt16Le(1)
        self.desktopSaveYGranularity = UInt16Le(20)
        self.pad2octetsA = UInt16Le(0)
        self.maximumOrderLevel = UInt16Le(1)
        self.numberFonts = UInt16Le()
        self.orderFlags = UInt16Le(OrderFlag.NEGOTIATEORDERSUPPORT)
        self.orderSupport = ArrayType(UInt8, init = [UInt8(0) for _ in range (0, 32)],  readLen = UInt8(32))
        self.textFlags = UInt16Le()
        self.orderSupportExFlags = UInt16Le()
        self.pad4octetsB = UInt32Le()
        self.desktopSaveSize = UInt32Le(480 * 480)
        self.pad2octetsC = UInt16Le()
        self.pad2octetsD = UInt16Le()
        self.textANSICodePage = UInt16Le(0)
        self.pad2octetsE = UInt16Le()
        
class BitmapCacheCapability(CompositeType):
    """
    Order use to cache bitmap very useful
    client -> server
    @see: http://msdn.microsoft.com/en-us/library/cc240559.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.pad1 = UInt32Le()
        self.pad2 = UInt32Le()
        self.pad3 = UInt32Le()
        self.pad4 = UInt32Le()
        self.pad5 = UInt32Le()
        self.pad6 = UInt32Le()
        self.cache0Entries = UInt16Le()
        self.cache0MaximumCellSize = UInt16Le()
        self.cache1Entries = UInt16Le()
        self.cache1MaximumCellSize = UInt16Le()
        self.cache2Entries = UInt16Le()
        self.cache2MaximumCellSize = UInt16Le()
        
class PointerCapability(CompositeType):
    """
    Use to indicate pointer handle of client
    Paint by server or per client
    client -> server
    server -> client
    @see: http://msdn.microsoft.com/en-us/library/cc240562.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.colorPointerFlag = UInt16Le()
        self.colorPointerCacheSize = UInt16Le()
        self.pointerCacheSize = UInt16Le()
        
class InputCapability(CompositeType):
    """
    Use to indicate input capabilities
    client -> server
    server -> client
    @see: http://msdn.microsoft.com/en-us/library/cc240563.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.inputFlags = UInt16Le()
        self.pad2octetsA = UInt16Le()
        #same value as gcc.ClientCoreSettings.kbdLayout
        self.keyboardLayout = UInt32Le()
        #same value as gcc.ClientCoreSettings.keyboardType
        self.keyboardType = UInt32Le()
        #same value as gcc.ClientCoreSettings.keyboardSubType
        self.keyboardSubType = UInt32Le()
        #same value as gcc.ClientCoreSettings.keyboardFnKeys
        self.keyboardFunctionKey = UInt32Le()
        #same value as gcc.ClientCoreSettingrrs.imeFileName
        self.imeFileName = String("\x00" * 64, readLen = UInt8(64))
        
class BrushCapability(CompositeType):
    """
    Use to indicate brush capability
    client -> server
    @see: http://msdn.microsoft.com/en-us/library/cc240564.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.brushSupportLevel = UInt32Le(BrushSupport.BRUSH_DEFAULT)
        
class GlyphCapability(CompositeType):
    """
    Use in font order
    client -> server
    @see: http://msdn.microsoft.com/en-us/library/cc240565.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.glyphCache = ArrayType(CacheEntry, init = [CacheEntry() for _ in range(0,10)], readLen = UInt8(10))
        self.fragCache = UInt32Le()
        #all fonts are sent with bitmap format (very expensive)
        self.glyphSupportLevel = UInt16Le(GlyphSupport.GLYPH_SUPPORT_NONE)
        self.pad2octets = UInt16Le()
        
class OffscreenBitmapCacheCapability(CompositeType):
    """
    use to cached bitmap in offscreen area
    client -> server
    @see: http://msdn.microsoft.com/en-us/library/cc240550.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.offscreenSupportLevel = UInt32Le(OffscreenSupportLevel.FALSE)
        self.offscreenCacheSize = UInt16Le()
        self.offscreenCacheEntries = UInt16Le()
        
class VirtualChannelCapability(CompositeType):
    """
    use to determine virtual channel compression
    client -> server
    server -> client
    @see: http://msdn.microsoft.com/en-us/library/cc240551.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.flags = UInt32Le(VirtualChannelCompressionFlag.VCCAPS_NO_COMPR)
        self.VCChunkSize = UInt32Le(optional = True)
        
class SoundCapability(CompositeType):
    """
    Use to exchange sound capability
    client -> server
    @see: http://msdn.microsoft.com/en-us/library/cc240552.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.soundFlags = UInt16Le(SoundFlag.NONE)
        self.pad2octetsA = UInt16Le()
        
class ControlCapability(CompositeType):
    """
    client -> server but server ignore contents! Thanks krosoft for brandwidth
    @see: http://msdn.microsoft.com/en-us/library/cc240568.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.controlFlags = UInt16Le()
        self.remoteDetachFlag = UInt16Le()
        self.controlInterest = UInt16Le(0x0002)
        self.detachInterest = UInt16Le(0x0002)
    
class WindowActivationCapability(CompositeType):
    """
    client -> server but server ignore contents! Thanks krosoft for brandwidth
    @see: http://msdn.microsoft.com/en-us/library/cc240569.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.helpKeyFlag = UInt16Le()
        self.helpKeyIndexFlag = UInt16Le()
        self.helpExtendedKeyFlag = UInt16Le()
        self.windowManagerKeyFlag = UInt16Le()
        
class FontCapability(CompositeType):
    """
    Use to indicate font support
    client -> server
    server -> client
    @see: http://msdn.microsoft.com/en-us/library/cc240571.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.fontSupportFlags = UInt16Le(0x0001)
        self.pad2octets = UInt16Le()
        
class ColorCacheCapability(CompositeType):
    """
    client -> server
    server -> client
    @see: http://msdn.microsoft.com/en-us/library/cc241564.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.colorTableCacheSize = UInt16Le(0x0006)
        self.pad2octets = UInt16Le()
        
class ShareCapability(CompositeType):
    """
    Use to advertise channel id of server
    client -> server
    server -> client
    @see: http://msdn.microsoft.com/en-us/library/cc240570.aspx
    """
    def __init__(self, readLen = None):
        CompositeType.__init__(self, readLen = readLen)
        self.nodeId = UInt16Le()
        self.pad2octets = UInt16Le()