Add chunked file transfer support

Author: bbarni2020

cloud/server.py

@@ -531,9 +531,55 @@ def handle_command_response(data):
     command_id = data.get('command_id')
     print(f'Received command_response for command_id: {command_id}')
     if command_id in pending_commands:
-        pending_commands[command_id]['response'] = data.get('response')
-        pending_commands[command_id]['completed'] = True
-        print(f'Updated pending_commands for command_id: {command_id}')
+        response_data = data.get('response')
+        if isinstance(response_data, dict) and response_data.get('chunked'):
+            print(f'Detected chunked response for command_id: {command_id}')
+            pending_commands[command_id]['chunked'] = True
+            pending_commands[command_id]['total_chunks'] = response_data.get('total_chunks', 0)
+            pending_commands[command_id]['total_size'] = response_data.get('total_size', 0)
+            pending_commands[command_id]['file_type'] = response_data.get('type', 'text')
+            pending_commands[command_id]['file_status'] = response_data.get('status', 'File content retrieved')
+            pending_commands[command_id]['chunks_received'] = {}
+            pending_commands[command_id]['completed'] = False
+        else:
+            pending_commands[command_id]['response'] = response_data
+            pending_commands[command_id]['completed'] = True
+            print(f'Updated pending_commands for command_id: {command_id}')
+
+@socketio.on('command_response_chunk')
+def handle_command_response_chunk(data):
+    command_id = data.get('command_id')
+    chunk_index = data.get('chunk_index')
+    total_chunks = data.get('total_chunks')
+    content = data.get('content')
+    is_final = data.get('is_final', False)
+
+    print(f'Received chunk {chunk_index + 1}/{total_chunks} for command_id: {command_id}')
+
+    if command_id in pending_commands:
+        command_data = pending_commands[command_id]
+        if 'chunks_received' not in command_data:
+            command_data['chunks_received'] = {}
+        command_data['chunks_received'][chunk_index] = content
+        if len(command_data['chunks_received']) == command_data.get('total_chunks', 0):
+            print(f'All chunks received for command_id: {command_id}, assembling response...')
+            assembled_content = ""
+            for i in range(command_data['total_chunks']):
+                if i in command_data['chunks_received']:
+                    assembled_content += command_data['chunks_received'][i]
+            final_response = {
+                'status': command_data.get('file_status', 'File content retrieved'),
+                'content': assembled_content,
+                'type': command_data.get('file_type', 'text')
+            }
+            command_data['response'] = final_response
+            command_data['completed'] = True
+            del command_data['chunks_received']
+            print(f'Assembled chunked response for command_id: {command_id}')
+        elif is_final:
+            print(f'Final chunk received but missing chunks for command_id: {command_id}')
+            command_data['response'] = {'error': 'Incomplete chunked file transfer'}
+            command_data['completed'] = True
 
 def send_command_to_server(server_id, command, command_id=None, method='GET', body=None, shareify_jwt=None):
     if server_id not in connected_servers:

current/cloud_connection.py

@@ -488,7 +488,14 @@ def execute_api_request(self, command_id, url, method='GET', body=None, shareify
 
             print(f"Making {method} request to: {full_url}")
 
-            if not (url == '/resources' or url == 'resources' or url == '/is_up' or url == 'is_up' or url == '/user/get_self' or url == 'user/get_self' or url == 'user/login' or url == '/user/login' or url == '/get_logs' or url == '/finder' or url == '/get_file'):
+            allowed_endpoints = [
+                '/resources', 'resources',
+                '/is_up', 'is_up',
+                '/user/get_self', 'user/get_self',
+                '/user/login', 'user/login',
+                '/get_logs', '/finder', '/get_file'
+            ]
+            if not any(url == ep or url.startswith(ep) for ep in allowed_endpoints):
                 if self.sio.connected:
                     try:
                         self.sio.emit('command_response', {
@@ -532,19 +539,12 @@ def execute_api_request(self, command_id, url, method='GET', body=None, shareify
                 response_data = response.json()
             except json.JSONDecodeError:
                 response_data = response.text
-            
-            if self.sio.connected:
-                try:
-                    self.sio.emit('command_response', {
-                        'command_id': command_id,
-                        'response': response_data
-                    })
-                    print(f"Successfully emitted response for command {command_id}")
-                    time.sleep(0.1)
-                except Exception as emit_error:
-                    print(f"Failed to emit response: {emit_error}")
+
+            is_file_endpoint = url.endswith('/get_file') or 'get_file' in url
+            if is_file_endpoint and isinstance(response_data, dict) and 'content' in response_data:
+                self.handle_large_file_response(command_id, response_data)
             else:
-                print("Socket not connected, cannot emit response")
+                self.send_standard_response(command_id, response_data)
 
         except requests.exceptions.Timeout:
             print(f"API request timeout for command {command_id}")

host/cloud_connection.py

@@ -540,18 +540,11 @@ def execute_api_request(self, command_id, url, method='GET', body=None, shareify
             except json.JSONDecodeError:
                 response_data = response.text
 
-            if self.sio.connected:
-                try:
-                    self.sio.emit('command_response', {
-                        'command_id': command_id,
-                        'response': response_data
-                    })
-                    print(f"Successfully emitted response for command {command_id}")
-                    time.sleep(0.1)
-                except Exception as emit_error:
-                    print(f"Failed to emit response: {emit_error}")
+            is_file_endpoint = url.endswith('/get_file') or 'get_file' in url
+            if is_file_endpoint and isinstance(response_data, dict) and 'content' in response_data:
+                self.handle_large_file_response(command_id, response_data)
             else:
-                print("Socket not connected, cannot emit response")
+                self.send_standard_response(command_id, response_data)
 
         except requests.exceptions.Timeout:
             print(f"API request timeout for command {command_id}")
@@ -585,6 +578,128 @@ def execute_api_request(self, command_id, url, method='GET', body=None, shareify
                     })
                 except Exception as emit_error:
                     print(f"Failed to emit general error: {emit_error}")
+
+    def send_standard_response(self, command_id, response_data):
+        if self.sio.connected:
+            try:
+                self.sio.emit('command_response', {
+                    'command_id': command_id,
+                    'response': response_data
+                })
+                print(f"Successfully emitted response for command {command_id}")
+                time.sleep(0.1)
+            except Exception as emit_error:
+                print(f"Failed to emit response: {emit_error}")
+        else:
+            print("Socket not connected, cannot emit response")
+
+    def handle_large_file_response(self, command_id, response_data):
+        import base64
+
+        if 'content' not in response_data:
+            self.send_standard_response(command_id, response_data)
+            return
+
+        content = response_data['content']
+
+        if response_data.get('type') == 'binary':
+            content_size = len(content) * 3 // 4
+        else:
+            content_size = len(content.encode('utf-8'))
+
+        chunk_threshold = 4 * 1024 * 1024
+
+        if content_size > chunk_threshold:
+            print(f"File content size ({content_size} bytes) exceeds threshold, chunking...")
+            self.send_chunked_file_response(command_id, response_data)
+        else:
+            print(f"File content size ({content_size} bytes) within limit, sending normally")
+            self.send_standard_response(command_id, response_data)
+
+    def send_chunked_file_response(self, command_id, response_data):
+        import base64
+
+        content = response_data['content']
+        file_type = response_data.get('type', 'text')
+        status = response_data.get('status', 'File content retrieved')
+
+        chunk_size = 3 * 1024 * 1024
+
+        if file_type == 'binary':
+            try:
+                original_bytes = base64.b64decode(content)
+                total_size = len(original_bytes)
+                chunks_needed = (total_size + chunk_size - 1) // chunk_size
+
+                print(f"Sending binary file in {chunks_needed} chunks...")
+
+                self.sio.emit('command_response', {
+                    'command_id': command_id,
+                    'response': {
+                        'status': status,
+                        'type': file_type,
+                        'chunked': True,
+                        'total_chunks': chunks_needed,
+                        'total_size': total_size,
+                        'chunk_index': 0
+                    }
+                })
+                time.sleep(0.1)
+
+                for i in range(chunks_needed):
+                    start_pos = i * chunk_size
+                    end_pos = min(start_pos + chunk_size, total_size)
+                    chunk_bytes = original_bytes[start_pos:end_pos]
+                    chunk_b64 = base64.b64encode(chunk_bytes).decode('utf-8')
+
+                    self.sio.emit('command_response_chunk', {
+                        'command_id': command_id,
+                        'chunk_index': i,
+                        'total_chunks': chunks_needed,
+                        'content': chunk_b64,
+                        'is_final': (i == chunks_needed - 1)
+                    })
+                    print(f"Sent chunk {i + 1}/{chunks_needed}")
+                    time.sleep(0.05)
+
+            except Exception as e:
+                print(f"Error chunking binary content: {e}")
+                self.send_standard_response(command_id, {'error': 'Failed to chunk binary content'})
+        else:
+            content_bytes = content.encode('utf-8')
+            total_size = len(content_bytes)
+            chunks_needed = (total_size + chunk_size - 1) // chunk_size
+
+            print(f"Sending text file in {chunks_needed} chunks...")
+
+            self.sio.emit('command_response', {
+                'command_id': command_id,
+                'response': {
+                    'status': status,
+                    'type': file_type,
+                    'chunked': True,
+                    'total_chunks': chunks_needed,
+                    'total_size': total_size,
+                    'chunk_index': 0
+                }
+            })
+            time.sleep(0.1)
+
+            for i in range(chunks_needed):
+                start_pos = i * chunk_size
+                end_pos = min(start_pos + chunk_size, total_size)
+                chunk_bytes = content_bytes[start_pos:end_pos]
+                chunk_text = chunk_bytes.decode('utf-8', errors='ignore')
+
+                self.sio.emit('command_response_chunk', {
+                    'command_id': command_id,
+                    'chunk_index': i,
+                    'total_chunks': chunks_needed,
+                    'content': chunk_text,
+                    'is_final': (i == chunks_needed - 1)
+                })
+                print(f"Sent chunk {i + 1}/{chunks_needed}")
+                time.sleep(0.05)
 
 def main():
     try: