Repository: agraubert/agutil

Issue:
SecureConnection.savefile() times out for large files
Right now the _timeout_ argument for SecureConnection.savefile() applies to the operation as a whole.  This is inconsequential for small files, but can timeout on larger files where the transfer may not complete in time even though the two sockets are still actively communicating.  The _timeout_ parameter should apply to each chunk of the file, so the operation will continue as long as it receives at least one chunk per timeout interval.

Discussion:

Patch:
diff --git a/README.md b/README.md
index 38e01c1..4e79ffb 100644
--- a/README.md
+++ b/README.md
@@ -34,5 +34,29 @@ The __security__ package:
 ## Documentation:
 Detailed documentation of these packages can be found on the [agutil Github wiki page](https://github.com/agraubert/agutil/wiki)
 
+## Features in development
+
+## agutil (main module)
+The following change has been made to the `agutil` module:
+* Added a `byteSize()` method to convert a number of bytes to a human-readable string
+
+##### API
+* byteSize(n):
+
+  Returns a string with _n_ converted to the highest unit of Bytes where _n_ would
+  be at least 1 (caps at ZiB), rounded to 1 decimal place.  Examples:
+
+  * byteSize(1) = `1B`
+  * byteSize(1023) = `1023B`
+  * byteSize(1024) = `1KiB`
+  * byteSize(856633344) = `816.9MiB`
+  * byteSize(12379856472314232739172) = `10.5ZiB`
+
+## agutil.security.SECURECONNECTION
+The following change has been made to the `agutil.security.SecureConnection` class:
+* `savefile()` timeout now applies to each chunk of the file.  The operation will
+block so long as the remote socket sends at least one chunk per timeout period.
+(API unchanged)
+
 ## Installation note:
 This package requires PyCrypto, which typically has issues compiling on windows.  If you are on windows and `pip install agutil` fails during the installation of PyCrypto, then follow the instructions [here](https://github.com/sfbahr/PyCrypto-Wheels) for installing PyCrypto from a precompiled wheel, and then run `pip install agutil` again.
diff --git a/agutil/__init__.py b/agutil/__init__.py
index 56dbd0a..58d133f 100644
--- a/agutil/__init__.py
+++ b/agutil/__init__.py
@@ -1,4 +1,4 @@
 from .src.search_range import search_range
 from .src.status_bar import status_bar
 from .src.logger import Logger, DummyLog
-from .src.misc import split_iterable, byte_xor, intToBytes, bytesToInt
+from .src.misc import split_iterable, byte_xor, intToBytes, bytesToInt, byteSize
diff --git a/agutil/security/src/connection.py b/agutil/security/src/connection.py
index ff825b9..2ecbd6e 100644
--- a/agutil/security/src/connection.py
+++ b/agutil/security/src/connection.py
@@ -1,5 +1,5 @@
 from .securesocket import SecureSocket
-from ... import io, Logger, DummyLog
+from ... import io, Logger, DummyLog, byteSize
 from . import protocols
 from socket import timeout as socketTimeout
 import threading
@@ -33,7 +33,7 @@ class SecureConnection:
         self.pending_tasks = threading.Event()
         self.intakeEvent = threading.Event()
         self.pendingRequest = threading.Event()
-        self.transferComplete = threading.Event()
+        self.transferTracking = {}
         self.killedtask = threading.Event()
         self.completed_transfers = set()
         self.queuedmessages = [] #Queue of decrypted received text messages
@@ -167,9 +167,10 @@ class SecureConnection:
             if not result:
                 raise socketTimeout("No file transfer requests recieved within the specified timeout")
             self.pendingRequest.clear()
-        (filename, auth) = self.authqueue.pop(0)
+        (filename, auth, size) = self.authqueue.pop(0)
         if not force:
             print("The remote socket is attempting to send the file '%s'"%filename)
+            print("Size:", byteSize(size))
             accepted = False
             choice = ""
             while not accepted:
@@ -189,19 +190,20 @@ class SecureConnection:
                 self.pending_tasks.set()
                 return
         self.log("Accepted transfer of file '%s'"%filename)
+        transferKey = '%s-%d'%(auth, hash(destination))
+        self.transferTracking[transferKey] = threading.Event()
         self.schedulingqueue.append({
             'cmd': protocols.lookupcmd('fti'),
             'auth': auth,
-            'filepath': destination
+            'filepath': destination,
+            'key': transferKey,
+            'timeout':timeout
         })
         self.pending_tasks.set()
         self.log("Waiting for transfer to complete...", "INFO")
-        while destination not in self.completed_transfers:
-            result = self.transferComplete.wait(timeout)
-            if not result:
-                raise socketTimeout("File transfer did not complete in the specified timeout")
-            self.transferComplete.clear()
-        self.completed_transfers.remove(destination)
+        self.transferTracking[transferKey].wait()
+        if destination not in self.completed_transfers:
+            raise socketTimeout("File transfer did not complete in the specified timeout")
         self.log("File transfer complete", "INFO")
         return destination
 
diff --git a/agutil/security/src/protocols.py b/agutil/security/src/protocols.py
index 58a512f..da79f11 100644
--- a/agutil/security/src/protocols.py
+++ b/agutil/security/src/protocols.py
@@ -1,6 +1,7 @@
 import rsa
 import os
 import random
+from socket import timeout as socketTimeout
 
 _COMMANDS = ['kill', 'ti', 'to', 'fri', 'fro', 'fto', 'fti', '_NULL', 'dci',]
 _CMD_LOOKUP = {}
@@ -150,7 +151,12 @@ def _file_request_out(sock,cmd,name):
     sock.filemap[auth_key] = cmd['filepath']
     sock.sock.sendAES(packcmd(
         'fri',
-        {'auth':auth_key, 'filename':os.path.basename(cmd['filepath']), 'name':name}
+        {
+            'auth':auth_key,
+            'filename':os.path.basename(cmd['filepath']),
+            'name':name,
+            'size': str(os.path.getsize(cmd['filepath']))
+        }
     ), '__cmd__')
     tasklog("Sent file transfer request to remote socket", "DETAIL")
     sock.schedulingqueue.append({
@@ -162,7 +168,7 @@ def _file_request_out(sock,cmd,name):
 def _file_request_in(sock,cmd,name):
     tasklog = sock.log.bindToSender(sock.log.name+":"+name)
     tasklog("Received new file request.  Queueing authorization", "DETAIL")
-    sock.authqueue.append((cmd['filename'], cmd['auth']))
+    sock.authqueue.append((cmd['filename'], cmd['auth'], int(cmd['size'])))
     sock.pendingRequest.set()
     sock.schedulingqueue.append({
         'cmd': lookupcmd('kill'),
@@ -196,12 +202,17 @@ def _file_transfer_in(sock,cmd,name):
         {'name':name, 'auth':cmd['auth'], 'reject':'reject' in cmd}
     ), '__cmd__')
     if 'reject' not in cmd:
-        sock.sock.recvRAW(name, timeout=None)
-        tasklog("File transfer initiated", "DEBUG")
-        sock.sock.recvAES(name, output_file=cmd['filepath'])
-        sock.completed_transfers.add(cmd['filepath'])
-        sock.transferComplete.set()
-        tasklog("Transfer complete", "DEBUG")
+        try:
+            sock.sock.recvRAW(name, timeout=cmd['timeout'])
+            tasklog("File transfer initiated", "DEBUG")
+            sock.sock.recvAES(name, output_file=cmd['filepath'], timeout=cmd['timeout'])
+        except socketTimeout:
+            tasklog("Transfer timed out", "ERROR")
+        else:
+            tasklog("Transfer complete", "DEBUG")
+            sock.completed_transfers.add(cmd['filepath'])
+        finally:
+            sock.transferTracking[cmd['key']].set()
     sock.schedulingqueue.append({
         'cmd': lookupcmd('kill'),
         'name': name
diff --git a/agutil/src/misc.py b/agutil/src/misc.py
index 28fefc6..0605346 100644
--- a/agutil/src/misc.py
+++ b/agutil/src/misc.py
@@ -1,4 +1,5 @@
 from itertools import chain, islice, zip_longest
+from math import log
 
 def intToBytes(num, padding_length=0):
     s = format(num, 'x')
@@ -28,3 +29,9 @@ def split_iterable(seq, length):
 
 def byte_xor(b1, b2):
     return intToBytes(bytesToInt(b1)^bytesToInt(b2), max(len(b1), len(b2)))
+
+def byteSize(n):
+    index = min(7, int(log(abs(n), 1024)))
+    suffix = ['B', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB', 'EiB', 'ZiB'][index]
+    output = '%.1f'%(n / 1024**index)
+    return output.rstrip('.0')+suffix


Test patch:
diff --git a/tests/test_misc_functions.py b/tests/test_misc_functions.py
index fce040a..1f96497 100644
--- a/tests/test_misc_functions.py
+++ b/tests/test_misc_functions.py
@@ -83,3 +83,14 @@ class test(unittest.TestCase):
             else:
                 self.assertEqual(len(converted), bytelen)
                 self.assertEqual(hash(converted), hash(bytestring))
+
+    def test_byte_size(self):
+        from agutil import byteSize
+        from math import log
+        import re
+        pattern = re.compile(r'(\d+(\.\d)?)([A-Z]i)?B')
+        for trial in range(25):
+            num = random.randint(0, sys.maxsize)
+            formatted = byteSize(num)
+            self.assertRegex(formatted, pattern)
+            self.assertLess(float(pattern.match(formatted).group(1)), 1024.0)