Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

porting remote support to smbprotocol #83

Merged
merged 2 commits into from
Apr 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
92 changes: 30 additions & 62 deletions ifcb/data/transfer/remote.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
Expand All @@ -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),
Expand Down
103 changes: 6 additions & 97 deletions ifcb/data/transfer/smb_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand All @@ -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
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading