porting remote support to smbprotocol

ifcb/data/transfer/remote.py

@@ -1,12 +1,12 @@
 import os
-
+import traceback
 from collections import defaultdict
 
-from .smb_utils import smb_connect, get_netbios_name, NameError
-from smb.base import SharedDevice
+from smbclient import smbclient
+
 
 DEFAULT_TIMEOUT = 30
-DEFAULT_SHARE = 'data'
+DEFAULT_SHARE = 'Data'
 
 class IfcbConnectionError(Exception):
     pass
@@ -15,68 +15,28 @@ def do_nothing(*args, **kw):
     pass
 
 class RemoteIfcb(object):
-    def __init__(self, addr, username, password, netbios_name=None, timeout=DEFAULT_TIMEOUT,
-                    share=DEFAULT_SHARE, directory='', connect=True):
+    def __init__(self, addr, username, password, timeout=DEFAULT_TIMEOUT,
+                    share=DEFAULT_SHARE, directory=''):
         self.addr = addr
         self.username = username
         self.password = password
         self.timeout = timeout
         self.share = share
-        self.connect = connect
-        self.netbios_name = netbios_name
         self.directory = directory
         self._c = None
     def open(self):
-        if self._c is not None:
-            return
-        try:
-            self._c = smb_connect(self.addr, self.username, self.password, self.netbios_name, self.timeout)
-        except:
-            raise IfcbConnectionError('unable to connect to IFCB')
+        smbclient.register_session(self.addr, self.username, self.password)
     def close(self):
-        if self._c is not None:
-            self._c.close()
-            self._c = None
+        smbclient.delete_session(self.addr)
     def __enter__(self):
-        if self.connect:
-            self.open()
+        self.open()
         return self
     def __exit__(self, type, value, traceback):
         self.close()
-    def ensure_connected(self):
-        if self._c is None:
-            raise IfcbConnectionError('IFCB is not connected')
-    def is_responding(self):
-        # tries to get NetBIOS name to see if IFCB is responding
-        if self.netbios_name is not None:
-            return True # FIXME determine connection state
-        if self._c is not None:
-            return True
-        else:
-            try:
-                get_netbios_name(self.addr, timeout=self.timeout)
-                return True
-            except:
-                return False
-    def list_shares(self):
-        self.ensure_connected()
-        for share in self._c.listShares():
-            if share.type == SharedDevice.DISK_TREE:
-                yield share.name
-    def share_exists(self):
-        self.ensure_connected()
-        for share in self.list_shares():
-            if share.lower() == self.share.lower():
-                return True
-        return False
     def list_filesets(self):
         """list fileset lids, most recent first"""
-        self.ensure_connected()
         fs = defaultdict(lambda: 0)
-        for f in self._c.listPath(self.share, self.directory):
-            if f.isDirectory:
-                continue
-            fn = f.filename
+        for fn in smbclient.listdir('\\'.join(['\\', self.addr, self.share, self.directory])):
             lid, ext = os.path.splitext(fn)
             if ext in ['.hdr','.roi','.adc']:
                 fs[lid] += 1
@@ -86,49 +46,57 @@ def list_filesets(self):
                 complete_sets.append(lid)
         return sorted(complete_sets, reverse=True)
     def transfer_fileset(self, lid, local_directory, skip_existing=True, create_directories=True):
-        self.ensure_connected()
         if create_directories:
             os.makedirs(local_directory, exist_ok=True)
         n_copied = 0
         for ext in ['hdr', 'adc', 'roi']:
             fn = '{}.{}'.format(lid, ext)
             local_path = os.path.join(local_directory, fn)
-            remote_path = os.path.join(self.directory, fn)
+            remote_path = '\\'.join(['\\', self.addr, self.share, self.directory, fn])
             temp_local_path = local_path + '.temp_download'
-
             if skip_existing and os.path.exists(local_path):
                 lf_size = os.path.getsize(local_path)
-                rf = self._c.getAttributes(self.share, remote_path)
-                if lf_size == rf.file_size:
+                stat = smbclient.stat(remote_path)
+                rf_size = stat.st_size
+                if lf_size == rf_size:
                     continue
 
-            with open(temp_local_path, 'wb') as fout:
-                self._c.retrieveFile(self.share, remote_path, fout, timeout=self.timeout)
+            with smbclient.open_file(remote_path, 'rb') as fin:
+                with open(temp_local_path, 'wb') as fout:
+                    while True:
+                        data = fin.read(1024)
+                        if not data:
+                            break
+                        fout.write(data)
+
             os.rename(temp_local_path, local_path)
             n_copied += 1
         return n_copied > 0
     def delete_fileset(self, lid):
-        self.ensure_connected()
-        for ext in ['hdr', 'adc', 'roi']:
-            self._c.deleteFiles(self.share, '{}.{}'.format(lid, ext))
+        #for ext in ['hdr', 'adc', 'roi']:
+        #    self._c.deleteFiles(self.share, '{}.{}'.format(lid, ext))
+        raise NotImplementedError()
     def sync(self, local_directory, progress_callback=do_nothing, fileset_callback=do_nothing):
         # local_directory can be
         # * a path, or
         # * a callbale returning a path when passed a bin lid
-        self.ensure_connected()
         fss = self.list_filesets()
         copied = []
         failed = []
         for lid in fss:
+            print(lid)
             try:
                 if callable(local_directory):
                     destination_directory = local_directory(lid)
+                else:
+                    destination_directory = local_directory
                 was_copied = self.transfer_fileset(lid, destination_directory, skip_existing=True)
                 if was_copied:
                     copied.append(lid)
                     fileset_callback(lid)
-            except:
+            except Exception as e:
                 failed.append(lid)
+                traceback.print_exc()
                 pass
             progress_callback({
                 'total': len(fss),

ifcb/data/transfer/smb_utils.py

@@ -4,8 +4,7 @@
 from contextlib import contextmanager
 import traceback
 
-from smb.SMBConnection import SMBConnection
-from nmb.NetBIOS import NetBIOS
+import smbclient
 
 DEFAULT_TIMEOUT=30
 
@@ -30,45 +29,22 @@ def path_on_share(path):
 def list_local_directory(path):
     return os.listdir(path)
 
-class NameError(Exception):
-    pass
-
-def get_netbios_name(remote_addr, timeout=DEFAULT_TIMEOUT):
-    nb = NetBIOS()
-    names = nb.queryIPForName(remote_addr, timeout=timeout)
-    nb.close()
-    if names is None or len(names) == 0:
-        raise NameError('No NetBIOS name found for {}'.format(remote_addr))
-    elif len(names) > 1:
-        logging.warn('More than one NetBIOS name for {}'.format(remote_addr))
-    return names[0]
 
-def smb_connect(remote_server, username, password, netbios_name=None, timeout=DEFAULT_TIMEOUT):
+def smb_connect(remote_server, username, password):
 
-    if netbios_name is None:
-        logging.debug('Querying NetBIOS for name of {}'.format(remote_server))
-        netbios_name = get_netbios_name(remote_server, timeout=timeout)
-        logging.debug('Name is {}'.format(netbios_name))
+    smbclient.register_session(remote_server, username, password)
 
-    logging.debug('Connecting to {}'.format(remote_server))
+    return smbclient
 
-    c = SMBConnection(username, password, 'ignore', netbios_name)
-    c.connect(remote_server, timeout=timeout)
-
-    return c
 
 @contextmanager
-def smb_connection(remote_server, username, password, netbios_name=None, timeout=DEFAULT_TIMEOUT):
-    c = smb_connect(remote_server, username, password, netbios_name, timeout)
+def smb_connection(remote_server, username, password):
+    c = smb_connect(remote_server, username, password)
 
     try:
         yield c
     except:
         traceback.print_exc()
-    finally:
-        logging.debug('Closing connection to {}'.format(remote_server))
-
-        c.close()
 
 def do_nothing(*args, **kw):
     pass
@@ -79,70 +55,3 @@ def progress(total_files, n_copied, fn):
         'copied': n_copied,
         'filename': fn,
     }
-
-def smb_sync_directory(remote_server, username, password, remote_path, local_path,
-    timeout=DEFAULT_TIMEOUT, limit=None, progress_callback=do_nothing):
-    """copies a remote directory to a local one (non-recursive).
-    remote_server = remote server DNS name or IP address
-    username = Samba username on remote server
-    password = Samba username's password on remote server
-    remote_path = path to remote directory including share name (e.g., '/some_share/some/folder')
-    local_path = path to local directory
-    will only transfer files that are 1) not in the local directory, or
-    2) are a different size than the one in the local directory.
-    Does not recurse into subdirectories."""
-    local_files = list_local_directory(local_path)
-
-    total_files = 0
-    n_copied = 0
-
-    with smb_connection(remote_server, username, password, timeout=timeout) as c:
-        logging.debug('listing remote directory {}'.format(remote_path))
-
-        share = share_name(remote_path)
-        pos = path_on_share(remote_path)
-        remote_files = [f for f in c.listPath(share, pos) if not f.isDirectory]
-
-        def safe_copy_file(remote_file, local_file):
-            logging.debug('Copying {} to {}'.format(remote_file.filename, local_file))
-            download_path = local_file + '.temp_download'
-            remote_path = os.path.join(pos, remote_file.filename)
-
-            try:
-                with open(download_path,'wb') as fout:
-                    c.retrieveFile(share, remote_path, fout)
-                os.rename(download_path, local_file)
-                return True
-            except:
-                return False
-            finally:
-                # clean up
-                if os.path.exists(download_path):
-                    os.remove(download_path)
-
-        for rf in remote_files:
-            total_files += 1
-
-            if limit is not None and n_copied == limit:
-                return
-
-            name = rf.filename
-
-            progress_callback(progress(total_files, n_copied, name))
-
-            local_file = os.path.join(local_path, name)
-            try:
-                if name not in local_files:
-                    logging.debug('{} does not exist locally'.format(name))
-                    if safe_copy_file(rf, local_file):
-                        n_copied += 1
-                else:
-                    remote_size = rf.file_size
-                    local_size = os.path.getsize(local_file)
-                    if local_size != remote_size:
-                        logging.debug('remote file {} is not the same size as local copy'.format(name))
-                        if safe_copy_file(rf, local_file):
-                            n_copied += 1
-            except:
-                traceback.print_exc()
-                # move on to next file

requirements.txt

@@ -6,4 +6,5 @@ Pillow==10.3.0
 rectpack==0.2.2
 scikit-image==0.22.0
 pysmb==1.2.9.1
+smbprotocol==1.13.0
 pyyaml==6.0.1