add DepotFile + improve CDNClient

* DepotFile for Manifest mappings
* CDNClient now handles content servers selection better
* CDNClient can auto get and cache cdn auth tokens and depot keys
* added get_chunk() to CDNClient, only ZIP compression, TODO LZMA

steam/client/builtins/apps.py

@@ -142,20 +142,20 @@ class Apps(object):
                                       timeout=10
                                       )
 
-    def get_depot_key(self, depot_id, app_id=0):
+    def get_depot_key(self, app_id, depot_id):
         """Get depot decryption key
 
-        :param depot_id: depot id
-        :type depot_id: :class:`int`
         :param app_id: app id
-        :type app_id: :class:`int`
+        :type  app_id: :class:`int`
+        :param depot_id: depot id
+        :type  depot_id: :class:`int`
         :return: `CMsgClientGetDepotDecryptionKeyResponse <https://github.com/ValvePython/steam/blob/39627fe883feeed2206016bacd92cf0e4580ead6/protobufs/steammessages_clientserver_2.proto#L533-L537>`_
         :rtype: proto message
         """
         return self.send_job_and_wait(MsgProto(EMsg.ClientGetDepotDecryptionKey),
                                       {
-                                       'depot_id': depot_id,
                                        'app_id': app_id,
+                                       'depot_id': depot_id,
                                       },
                                       timeout=10
                                       )
@@ -182,9 +182,9 @@ class Apps(object):
         """Get access tokens
 
         :param app_ids: list of app ids
-        :type app_ids: :class:`list`
+        :type  app_ids: :class:`list`
         :param package_ids: list of package ids
-        :type package_ids: :class:`list`
+        :type  package_ids: :class:`list`
         :return: dict with ``apps`` and ``packages`` containing their access tokens, see example below
         :rtype: :class:`dict`, :class:`None`
 

steam/client/cdn.py

@@ -2,9 +2,11 @@
 from collections import OrderedDict, deque
 from six import itervalues
 import vdf
+
 from steam import webapi
-from steam.enums import EServerType
+from steam.enums import EResult, EServerType
 from steam.util.web import make_requests_session
+from stema.core.crypto symmetric_decrypt
 from steam.core.manifest import DepotManifest
 
 
@@ -67,12 +69,14 @@ class CDNClient(object):
         self.app_id = app_id
         self.web = make_requests_session()
         self.servers = deque()
+        self.cdn_auth_tokens = {}
+        self.depot_keys = {}
 
     @property
     def cell_id(self):
         return self.steam.cell_id
 
-    def init_servers(self, num_servers=10):
+    def init_servers(self, num_servers=20):
         self.servers.clear()
 
         for ip, port in self.steam.servers[EServerType.CS]:
@@ -85,32 +89,85 @@ class CDNClient(object):
         if not self.servers:
             raise RuntimeError("No content servers on SteamClient instance. Is it logged in?")
 
-    def get_content_server(self):
-        server = self.servers[0]
-        self.servers.rotate(-1)
-        return server
+    def get_content_server(self, rotate=True):
+        if rotate:
+            self.servers.rotate(-1)
+        return self.servers[0]
+
+    def get_cdn_auth_token(self, depot_id):
+        if depot_id not in self.cdn_auth_tokens:
+            msg = self.steam.get_cdn_auth_token(depot_id, 'steampipe.steamcontent.com')
+
+            if msg.eresult == EResult.OK:
+                self.cdn_auth_tokens[depot_id] = msg.token
+            elif msg is None:
+                raise Exception("Failed getting depot key: %s" % repr(EResult.Timeout))
+            else:
+                raise Exception("Failed getting depot key: %s" % repr(EResult(msg.eresult)))
+
+        return self.cdn_auth_tokens[depot_id]
+
+    def get_depot_key(self, depot_id):
+        if depot_id not in self.depot_keys:
+            msg = self.steam.get_depot_key(self.app_id, depot_id)
+            if msg.eresult == EResult.OK:
+                self.depot_keys[depot_id] = msg.depot_encryption_key
+            elif msg is None:
+                raise Exception("Failed getting depot key: %s" % repr(EResult.Timeout))
+            else:
+                raise Exception("Failed getting depot key: %s" % repr(EResult(msg.eresult)))
+
+        return self.depot_keys[depot_id]
 
     def get(self, command, args, auth_token=''):
         server = self.get_content_server()
 
-        url = "%s://%s:%s/%s/%s%s" % (
-            'https' if server.https else 'http',
-            server.host,
-            server.port,
-            command,
-            args,
-            auth_token,
-            )
+        while True:
+            url = "%s://%s:%s/%s/%s%s" % (
+                'https' if server.https else 'http',
+                server.host,
+                server.port,
+                command,
+                args,
+                auth_token,
+                )
+            resp = self.web.get(url)
 
-        return self.web.get(url)
+            if resp.ok:
+                return resp
+            elif resp.status_code in (401, 403, 404):
+                resp.raise_for_status()
 
-    def get_manifest(self, depot_id, manifest_id, auth_token):
-        resp = self.get('depot', '%s/manifest/%s/5' % (depot_id, manifest_id), auth_token)
+            server = self.get_content_server(rotate=True)
 
-        resp.raise_for_status()
+    def get_manifest(self, depot_id, manifest_id, cdn_auth_token=None, decrypt=True):
+        if cdn_auth_token is None:
+            cdn_auth_token = self.get_cdn_auth_token(depot_id)
+
+        resp = self.get('depot', '%s/manifest/%s/5' % (depot_id, manifest_id), cdn_auth_token)
 
         if resp.ok:
-            return DepotManifest(resp.content)
+            manifest = DepotManifest(resp.content)
+            if decrypt:
+                manifest.decrypt_filenames(self.get_depot_key(depot_id))
+            return manifest
+
+    def get_chunk(self, depot_id, chunk_id, cdn_auth_token=None):
+        if cdn_auth_token is None:
+            cdn_auth_token = self.get_cdn_auth_token(depot_id)
+
+        resp = self.get('depot', '%s/chunk/%s' % (depot_id, chunk_id), cdn_auth_token)
+
+        if resp.ok:
+            data = symmetric_decrypt(resp.content, self.get_depot_key(depot_id))
+
+            if data[:2] == b'VZ':
+                raise Exception("Implement LZMA lol")
+            else:
+                with ZipFile(BytesIO(data)) as zf:
+                    data = zf.read(zf.filelist[0])
+
+            return data
 
 
 class ContentServer(object):

steam/core/manifest.py

@@ -1,10 +1,12 @@
 
 from base64 import b64decode
 from io import BytesIO
-from zipfile import ZipFile, ZIP_DEFLATED
+from zipfile import ZipFile, ZIP_DEFLATED, BadZipFile
 from struct import pack
 from datetime import datetime
+from fnmatch import fnmatch
 
+from steam.enums import EDepotFileFlag
 from steam.core.crypto import symmetric_decrypt
 from steam.util.binary import StructReader
 from steam.protobufs.content_manifest_pb2 import (ContentManifestMetadata,
@@ -18,7 +20,12 @@ class DepotManifest(object):
     PROTOBUF_SIGNATURE_MAGIC = 0x1B81B817
     PROTOBUF_ENDOFMANIFEST_MAGIC = 0x32C415AB
 
-    def __init__(self, data):
+    def __init__(self, data=None):
+        """Manage depot manifest
+
+        :param data: manifest data
+        :type  data: bytes
+        """
         self.metadata = ContentManifestMetadata()
         self.payload = ContentManifestPayload()
         self.signature = ContentManifestSignature()
@@ -28,8 +35,8 @@ class DepotManifest(object):
 
     def __repr__(self):
         params = ', '.join([
-                    str(self.metadata.depot_id),
-                    str(self.metadata.gid_manifest),
+                    str(self.depot_id),
+                    str(self.gid),
                     repr(datetime.utcfromtimestamp(self.metadata.creation_time).isoformat().replace('T', ' ')),
                     ])
 
@@ -41,9 +48,35 @@ class DepotManifest(object):
             params,
             )
 
+    @property
+    def depot_id(self):
+        return self.metadata.depot_id
+
+    @property
+    def gid(self):
+        return self.metadata.gid_manifest
+
+    @property
+    def creation_time(self):
+        return self.metadata.creation_time
+
+    @property
+    def size_original(self):
+        return self.metadata.cb_disk_original
+
+    @property
+    def size_compressed(self):
+        return self.metadata.cb_disk_compressed
+
     def decrypt_filenames(self, depot_key):
+        """Decrypt all filenames in the manifest
+
+        :param depot_key: depot key
+        :type  depot_key: bytes
+        :raises: :class:`RuntimeError`
+        """
         if not self.metadata.filenames_encrypted:
-            return True
+            return
 
         for mapping in self.payload.mappings:
             filename = b64decode(mapping.filename)
@@ -51,17 +84,25 @@ class DepotManifest(object):
             try:
                 filename = symmetric_decrypt(filename, depot_key)
             except Exception:
-                print("Unable to decrypt filename for depot manifest")
-                return False
+                RuntimeError("Unable to decrypt filename for depot manifest")
 
             mapping.filename = filename
 
         self.metadata.filenames_encrypted = False
-        return True
 
     def deserialize(self, data):
-        with ZipFile(BytesIO(data)) as zf:
-            data = StructReader(zf.read(zf.filelist[0]))
+        """Deserialize a manifest (compressed or uncompressed)
+
+        :param data: manifest data
+        :type  data: bytes
+        """
+        try:
+            with ZipFile(BytesIO(data)) as zf:
+                data = zf.read(zf.filelist[0])
+        except BadZipFile:
+            pass
+
+        data = StructReader(data)
 
         magic, length = data.unpack('<II')
 
@@ -92,7 +133,12 @@ class DepotManifest(object):
         if magic != DepotManifest.PROTOBUF_ENDOFMANIFEST_MAGIC:
             raise Exception("Expecting end of manifest")
 
-    def serialize(self):
+    def serialize(self, compress=True):
+        """Serialize manifest
+
+        :param compress: wether the output should be Zip compressed
+        :type  compress: bytes
+        """
         data = BytesIO()
 
         part = self.payload.SerializeToString()
@@ -109,8 +155,79 @@ class DepotManifest(object):
 
         data.write(pack('<I', DepotManifest.PROTOBUF_ENDOFMANIFEST_MAGIC))
 
-        zbuff = BytesIO()
-        with ZipFile(zbuff, 'w', ZIP_DEFLATED) as zf:
-            zf.writestr('z', data.getvalue())
+        if compress:
+            zbuff = BytesIO()
+            with ZipFile(zbuff, 'w', ZIP_DEFLATED) as zf:
+                zf.writestr('z', data.getvalue())
+
+            return zbuff.getvalue()
+        else:
+            return data.getvalue()
+
+    def __iter__(self):
+        for mapping in self.payload.mappings:
+            yield DepotFile(self, mapping)
+
+    def iter_files(self, pattern=None):
+        """
+        :param pattern: unix shell wildcard pattern, see :module:`.fnmatch`
+        :type  pattern: str
+        """
+        for mapping in self.payload.mappings:
+            if (pattern is not None
+               and not fnmatch(mapping.filename.rstrip('\x00 \n\t'), pattern)):
+                continue
+            yield DepotFile(self, mapping)
+
+    def __len__(self):
+        return len(self.payload.mappings)
+
+
+class DepotFile(object):
+    def __init__(self, manifest, file_mapping):
+        """Depot file
+        :param manifest: depot manifest
+        :type  manifest: :class:`.DepotManifest`
+        :param file_mapping: depot file mapping instance
+        :type  file_mapping: ContentManifestPayload.FileMapping
+        """
+        if not isinstance(manifest, DepotManifest):
+            raise ValueError("Expected 'manifest' to be of type DepotManifest")
+        if not isinstance(file_mapping, ContentManifestPayload.FileMapping):
+            raise ValueError("Expected 'file_mapping' to be of type ContentManifestPayload.FileMapping")
+
+        self.manifest = manifest
+        self.file_mapping = file_mapping
+
+    def __repr__(self):
+        return "<%s(%s, %s, %s, %s)>" % (
+            self.__class__.__name__,
+            self.manifest.depot_id,
+            self.manifest.gid,
+            repr(self.filename),
+            'is_directory=True' if self.is_directory else self.size,
+            )
+
+    @property
+    def filename(self):
+        return self.file_mapping.filename.rstrip('\x00 \n\t')
+
+    @property
+    def size(self):
+        return self.file_mapping.size
+
+    @property
+    def chunks(self):
+        return self.file_mapping.chunks
+
+    @property
+    def flags(self):
+        return self.file_mapping.flags
+
+    @property
+    def is_directory(self):
+        return self.flags & EDepotFileFlag.Directory > 0
 
-        return zbuff.getvalue()
+    @property
+    def is_file(self):
+        return not self.is_directory

steam/enums/common.py

@@ -539,6 +539,19 @@ class ECurrencyCode(SteamIntEnum):
     Max = 42
 
 
+class EDepotFileFlag(SteamIntEnum):
+    UserConfig = 1
+    VersionedUserConfig = 2
+    Encrypted = 4
+    ReadOnly = 8
+    Hidden = 16
+    Executable = 32
+    Directory = 64
+    CustomExecutable = 128
+    InstallScript = 256
+    Symlink = 512
+
+
 # Do not remove
 from enum import EnumMeta