You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
shepherd-agent/shepherd/modules/uploader.py

211 lines
8.2 KiB

import shutil
import os
import threading
import paramiko
import shepherd.config
import shepherd.module
# configdef = shepherd.config.definition()
# Can either import shepherd.config here, and call a function to build a config_def
# or can leave a config_def entry point.
# probably go with entry point, to stay consistent with the module
class UploaderConfDef(shepherd.config.ConfDefinition):
def __init__(self):
super().__init__()
dests = self.add_def('destination', shepherd.config.TableArrayDef())
dests.add_def('name', shepherd.config.StringDef())
dests.add_def('protocol', shepherd.config.StringDef())
dests.add_def('address', shepherd.config.StringDef(optional=True))
dests.add_def('port', shepherd.config.IntDef(optional=True))
dests.add_def('path', shepherd.config.StringDef(optional=True))
dests.add_def('username', shepherd.config.StringDef(optional=True))
dests.add_def('password', shepherd.config.StringDef(optional=True))
dests.add_def('keyfile', shepherd.config.StringDef(default="", optional=True))
dests.add_def('add_id_to_path', shepherd.config.BoolDef(default=True, optional=True))
buckets = self.add_def('bucket', shepherd.config.TableArrayDef())
buckets.add_def('name', shepherd.config.StringDef())
buckets.add_def('open_link_on_new', shepherd.config.BoolDef())
buckets.add_def('opportunistic', shepherd.config.BoolDef(default=True, optional=True))
buckets.add_def('keep_copy', shepherd.config.BoolDef())
buckets.add_def('destination', shepherd.config.StringDef())
# on server side, we want to be able to list commands that a module responds to
# without actually instantiating the module class. Add command templates into
# the conf_def, than attach to them in the interface? Was worried about having
# "two sources of truth", but you already need to match the conf_def to the
# name where you access the value in the module. Could have add_command, which
# you then add standard conf_def subclasses to, to reuse validation and server
# form generation logic...
class UploaderInterface(shepherd.module.Interface):
def __init__(self, module):
super().__init__(module)
# self.add_command("trigger", self.module.camera_job)
def _get_bucket_path(self, bucket_name):
return self._module.buckets[bucket_name].path
def move_to_bucket(self, filepath, bucket_name):
dest_path = os.path.join(self._get_bucket_path(bucket_name),
os.path.basename(filepath))
temp_dest_path = dest_path + ".writing"
shutil.move(filepath, temp_dest_path)
os.rename(temp_dest_path, dest_path)
# notify bucket to check for new files
self._module.buckets[bucket_name].newfile_event.set()
class Destination():
def __init__(self, config, core_interface):
self.config = config
self.shepherd = core_interface
self.sendlist_condition = threading.Condition()
self.send_list = []
self.thread = threading.Thread(target=self._send_files)
self.thread.start()
# Override this in subclasses, implementing the actual upload process.
# Return true on success, false on failure.
def upload(self, filepath, suffix):
print ("Dummy uploading "+filepath)
return True
def add_files_to_send(self, file_path_list):
self.sendlist_condition.acquire()
for file_path in file_path_list:
if file_path not in self.send_list:
self.send_list.append(file_path)
self.sendlist_condition.notify()
self.sendlist_condition.release()
def _file_available(self):
return len(self.send_list) > 0
def _send_files(self):
while True:
self.sendlist_condition.acquire()
# this drops through immediately if there is something to send, otherwise waits
self.sendlist_condition.wait_for(self._file_available)
file_to_send = self.send_list.pop(0)
os.rename(file_to_send, file_to_send+".uploading")
self.sendlist_condition.release()
# Rename uploaded file to end with ".uploaded" on success, or back
# to original path on failure.
try:
self.upload(file_to_send, ".uploading")
os.rename(file_to_send+".uploading", file_to_send+".uploaded")
except:
os.rename(file_to_send+".uploading", file_to_send)
self.send_list.append(file_to_send)
class SFTPDestination(Destination):
def upload(self, filepath, suffix):
with paramiko.Transport((self.config["address"],
self.config["port"])) as transport:
transport.connect(username=self.config["username"],
password=self.config["password"])
with paramiko.SFTPClient.from_transport(transport) as sftp:
print("Uploading "+filepath+" to "+self.config["address"]+" via SFTP")
if self.config["add_id_to_path"]:
destdir = os.path.join(self.config["path"],
self.shepherd.id)
else:
destdir = self.config["path"]
try:
sftp.listdir(destdir)
except IOError:
print("Creating remot dir:" + destdir)
sftp.mkdir(destdir)
print("Target dir:"+destdir)
sftp.put(filepath+suffix,
os.path.join(destdir, os.path.basename(filepath)))
class Bucket():
def __init__(self, name, open_link_on_new, opportunistic, keep_copy,
destination, core_interface, path=None, old_path=None):
self.newfile_event = threading.Event()
self.newfile_event.set()
self.destination = destination
self.shepherd = core_interface
self.path = path
if self.path is None:
self.path = os.path.join(self.shepherd.root_dir, name)
if not os.path.exists(self.path):
os.makedirs(self.path)
if keep_copy:
self.old_path = old_path
if self.old_path is None:
self.old_path = os.path.join(self.shepherd.root_dir, name + "_old")
if not os.path.exists(self.old_path):
os.makedirs(self.old_path)
self.thread = threading.Thread(target=self._check_files)
self.thread.start()
def _check_files(self):
while True:
self.newfile_event.wait(timeout=10)
self.newfile_event.clear()
bucket_files = []
for item in os.listdir(self.path):
item_path = os.path.join(self.path, item)
if (os.path.isfile(item_path) and
(not item.endswith(".writing")) and
(not item.endswith(".uploading")) and
(not item.endswith(".uploaded"))):
bucket_files.append(item_path)
if bucket_files:
self.destination.add_files_to_send(bucket_files)
class UploaderModule(shepherd.module.Module):
conf_def = UploaderConfDef()
def __init__(self, config, core_interface):
super().__init__(config, core_interface)
print("Uploader config:")
print(self.config)
self.interface = UploaderInterface(self)
self.destinations = {}
self.buckets = {}
for dest_conf in self.config["destination"]:
if dest_conf["protocol"] == "sftp":
self.destinations[dest_conf["name"]] = SFTPDestination(dest_conf, core_interface)
else:
self.destinations[dest_conf["name"]] = Destination(dest_conf, core_interface)
for bucketconf in self.config["bucket"]:
bucketconf["destination"] = self.destinations[bucketconf["destination"]]
self.buckets[bucketconf["name"]] = Bucket(
**bucketconf, core_interface=self.shepherd)
def init_other_modules(self, interfaces): # pylint: disable=W0235
super().init_other_modules(interfaces)
if __name__ == "__main__":
pass
#print("main")
#main(sys.argv[1:])