445 lines
19 KiB
Python
445 lines
19 KiB
Python
#!/usr/bin/python3
|
||
|
||
import sys, hashlib, os, sqlite3, shutil, json
|
||
|
||
################################################################################
|
||
# program wide constants
|
||
################################################################################
|
||
|
||
GENERAL_ERROR=1
|
||
USAGE_ERROR=2
|
||
PATH_ERROR=3
|
||
|
||
################################################################################
|
||
# functions
|
||
################################################################################
|
||
|
||
def calculate_file_hash(algorithm, file_path, segment_size=4096):
|
||
file_descriptor = open(file_path, "rb")
|
||
hash_object = hashlib.new(algorithm)
|
||
segment = file_descriptor.read(segment_size)
|
||
while not segment==b'':
|
||
hash_object.update(segment)
|
||
segment = file_descriptor.read(segment_size)
|
||
return hash_object.hexdigest()
|
||
|
||
def calculate_parity(byte_string):
|
||
parity_byte = 0b00000000
|
||
for i in range(len(byte_string)):
|
||
parity_byte = parity_byte ^ byte_string[i]
|
||
return parity_byte
|
||
|
||
def gegerate_parity_file(input_path, parity_bytes, output_path):
|
||
input_file = open(input_path, "rb")
|
||
output_file = open(output_path, "wb")
|
||
|
||
segment = input_file.read(parity_bytes)
|
||
while not segment==b'':
|
||
output_file.write((calculate_parity(segment)).to_bytes(1, byteorder='big'))
|
||
segment = input_file.read(parity_bytes)
|
||
|
||
input_file.close()
|
||
output_file.close()
|
||
|
||
def create_container(storage_directory, parity=False, parity_bytes=512, checksum_algorithm='sha512', compress=False):
|
||
# prepare storage directory
|
||
os.makedirs(storage_directory)
|
||
os.mkdir(os.path.join(storage_directory, "objects"))
|
||
if (parity):
|
||
os.mkdir(os.path.join(storage_directory, "parity"))
|
||
|
||
db = sqlite3.connect(os.path.join(storage_directory, "container.sqlite"))
|
||
cursor = db.cursor()
|
||
# settings
|
||
cursor.execute("CREATE TABLE settings (option TEXT, value TEXT);")
|
||
cursor.execute("INSERT INTO settings VALUES ('parity', ?);", (str(parity),))
|
||
cursor.execute("INSERT INTO settings VALUES ('parity_bytes', ?);", (str(parity_bytes),))
|
||
cursor.execute("INSERT INTO settings VALUES ('checksum_algorithm', ?);", (checksum_algorithm,))
|
||
cursor.execute("INSERT INTO settings VALUES ('compress', ?);", (str(compress),))
|
||
# container
|
||
cursor.execute("CREATE TABLE hashes (id INTEGER PRIMARY KEY, hash TEXT UNIQUE);") # primary key somehow makes this similar to auto-increment (not actually auto increment but good enough for my use)
|
||
cursor.execute("CREATE TABLE tags (id INTEGER, tag TEXT);")
|
||
db.commit()
|
||
db.close()
|
||
|
||
def add_tag(storage_directory, file_hash, tag):
|
||
db = sqlite3.connect(os.path.join(storage_directory, "container.sqlite"))
|
||
cursor = db.cursor()
|
||
|
||
cursor.execute("SELECT id FROM hashes WHERE hash=?", (file_hash,))
|
||
row = cursor.fetchone()
|
||
if row == None:
|
||
cursor.execute("INSERT INTO hashes (hash) VALUES (?)", (file_hash,))
|
||
db.commit()
|
||
cursor.execute("SELECT id FROM hashes WHERE hash=?", (file_hash,))
|
||
row = cursor.fetchone()
|
||
internal_id = row[0]
|
||
|
||
cursor.execute("SELECT tag FROM tags WHERE id=?", (internal_id,))
|
||
rows = cursor.fetchall()
|
||
tags = []
|
||
for row in rows:
|
||
tags = tags+[row[0]]
|
||
|
||
tag_already_present = True
|
||
if not tag in tags:
|
||
tag_already_present = False
|
||
cursor.execute("INSERT INTO tags (id, tag) VALUES (?, ?)", (internal_id, tag))
|
||
db.commit()
|
||
|
||
db.close()
|
||
return not tag_already_present
|
||
|
||
def get_tags_by_hash(storage_directory, file_hash):
|
||
db = sqlite3.connect(os.path.join(storage_directory, "container.sqlite"))
|
||
cursor = db.cursor()
|
||
|
||
tags = None
|
||
cursor.execute("SELECT id FROM hashes WHERE hash=?", (file_hash,))
|
||
row = cursor.fetchone()
|
||
if not row == None:
|
||
internal_id = row[0]
|
||
cursor.execute("SELECT tag FROM tags WHERE id=?", (internal_id,))
|
||
rows = cursor.fetchall()
|
||
tags = []
|
||
for row in rows:
|
||
tags = tags+[row[0]]
|
||
|
||
db.close()
|
||
return tags
|
||
|
||
def get_hashes_by_tag(storage_directory, tag):
|
||
db = sqlite3.connect(os.path.join(storage_directory, "container.sqlite"))
|
||
cursor = db.cursor()
|
||
|
||
cursor.execute("SELECT id FROM tags WHERE tag=?", (tag,))
|
||
rows = cursor.fetchall()
|
||
internal_ids = []
|
||
for row in rows:
|
||
internal_ids = internal_ids+[row[0]]
|
||
|
||
file_hashes = []
|
||
for internal_id in internal_ids:
|
||
cursor.execute("SELECT hash FROM hashes WHERE id=?", (internal_id,))
|
||
row = cursor.fetchone()
|
||
file_hashes = file_hashes+[row[0]]
|
||
|
||
return file_hashes
|
||
|
||
def file_hash_or_path_is_known_hash(storage_directory, file_hash_or_path, compress):
|
||
suffix=""
|
||
if compress:
|
||
suffix=".xz"
|
||
#that last part is needed bc os.path.join automatically discards everything before the last element containing an absolute path so if file_hash_or_path is an absolute path that would be what join gives back
|
||
return os.path.isfile(os.path.join(storage_directory, "objects", file_hash_or_path+suffix)) and not os.path.isfile(file_hash_or_path)
|
||
|
||
def load_container_settings(storage_directory):
|
||
if not os.path.isfile(os.path.join(storage_directory, "container.sqlite")):
|
||
return (PATH_ERROR, None, None, None, None)
|
||
|
||
db = sqlite3.connect(os.path.join(storage_directory, "container.sqlite"))
|
||
cursor = db.cursor()
|
||
cursor.execute("SELECT option, value FROM settings")
|
||
|
||
#TODO: check container settings properly instead of just assuming default values for things that aren't found
|
||
# set default values and then read the db, just in case...
|
||
parity = False
|
||
parity_bytes = 512
|
||
checksum_algorithm = "sha512"
|
||
compress = False
|
||
for row in cursor.fetchall():
|
||
if row[0]=="parity":
|
||
parity = row[1]==str(True)
|
||
if row[0]=="parity_bytes":
|
||
parity_bytes = int(row[1])
|
||
if row[0]=="checksum_algorithm":
|
||
checksum_algorithm = row[1]
|
||
if row[0]=="compress":
|
||
compress = row[1]==str(True)
|
||
db.close()
|
||
|
||
# check storage container against settings
|
||
#TODO: check compression ?
|
||
if parity and not os.path.isdir(os.path.join(storage_directory, "parity")):
|
||
return (GENERAL_ERROR, None, None, checksum_algorithm, compress)
|
||
return (0, parity, parity_bytes, checksum_algorithm, compress)
|
||
|
||
################################################################################
|
||
# main program
|
||
################################################################################
|
||
|
||
if __name__ == "__main__":
|
||
|
||
USAGE="""Usage:
|
||
"""+sys.argv[0]+""" create <storage directory> [parity=<on|off>] [parity-bytes=<number of bytes for each parity byte>] [checksum-algorithm=<algorithm>] [compress=<on|off>] - set up a new storage directory
|
||
"""+sys.argv[0]+""" add <storage directory> <hash|file> <tags ...> - add tags to a file in the storage, if file is not already in the storage, same as add+file
|
||
"""+sys.argv[0]+""" add+file <storage directory> <file> <tags ...> - copy a file to the storage and add tags
|
||
"""+sys.argv[0]+""" add+file+move <storage directory> <file> <tags ...> - move a file to the storage and add tags
|
||
"""+sys.argv[0]+""" remove <storage directory> <hash|unique tag or tag set> - remove a file from the storage, return error if not found or multiple found
|
||
"""+sys.argv[0]+""" remove+multi <storage directory> <exact tag or set of exact tags> - remove a ll found files from the storage, return error if not found
|
||
"""+sys.argv[0]+""" search <storage directory> <tags or partial tags> - return paths and tags of all found files
|
||
"""+sys.argv[0]+""" search+first <storage directory> <tags or partial tags> - return hash and tags of first found file
|
||
"""+sys.argv[0]+""" search+unique <storage directory> <tags or partial tags> - return hash and tags of the found file, return error if not found or multiple found
|
||
"""+sys.argv[0]+""" lookup <storage directory> <hash|exact tag|set of exact tags> - return hash and tags of all found files
|
||
"""+sys.argv[0]+""" lookup+first <storage directory> <hash|exact tag|set of exact tags> - return hash and tags of first found file
|
||
"""+sys.argv[0]+""" lookup+unique <storage directory> <hash|exact tag|set of exact tags> - return hash and tags of the found file, return error if not found or multiple found
|
||
"""+sys.argv[0]+""" link <storage directory> <hash> <location> - add a symlink in <location> that points to the referenced file
|
||
"""+sys.argv[0]+""" check <storage directory> <hash> - check file contents against hash
|
||
"""+sys.argv[0]+""" check+parity <storage directory> <hash> - check file contents against hash and parity file
|
||
"""+sys.argv[0]+""" check+all <storage directory> - check all files against their hashes
|
||
"""+sys.argv[0]+""" check+all+parity <storage directory> - check all files against their hashes and parity files
|
||
"""+sys.argv[0]+""" update <storage directory> <hash> - update the hash (and parity if applickable) of the specified file (specify by previous hash)
|
||
"""+sys.argv[0]+""" update+all <storage directory> - update hashes (and parities if applickable) of all mismatching files
|
||
"""+sys.argv[0]+""" fix <storage directory> <hash> - attempt to fix the file using parity
|
||
"""+sys.argv[0]+""" fix+all <storage directory> - attempt to fix all files using parity
|
||
"""+sys.argv[0]+""" help - display this message
|
||
"""
|
||
#TODO: +path modifier for things that return a hash to return the path to the stored file instead
|
||
#TODO: +hash and +tags modifier for lookup
|
||
#TODO: condense modifiers onto the same lines as the main subcommand where possible
|
||
#TODO: clarification of <> and []
|
||
#TODO: subcommand to change container settings
|
||
VALID_COMMANDS=["create", "add", "remove", "search", "lookup", "link", "check", "update", "fix", "help"]
|
||
#TODO: (*fully) implemented subcommands: *create, *add, *lookup, *link, *help
|
||
#TODO: unimplemented subcommands: remove, search, check, update, fix
|
||
|
||
try:
|
||
command = sys.argv[1].split("+")
|
||
except IndexError:
|
||
print("No subcommand specified.", file=sys.stderr)
|
||
print(USAGE, file=sys.stderr)
|
||
sys.exit(USAGE_ERROR)
|
||
if not command[0] in VALID_COMMANDS:
|
||
print("Invalid command: "+command[0], file=sys.stderr)
|
||
print(USAGE, file=sys.stderr)
|
||
sys.exit(USAGE_ERROR)
|
||
|
||
# create subcommand: create a new directory containing a folder for stored objects, one for parity files and one for
|
||
# arguments: <storage directory> [parity=<on|off>] [parity-bytes=<number of bytes for each parity byte>] [checksum-algorithm=<algorithm>] [compress=<on|off>]
|
||
if command[0] == "create":
|
||
if len(sys.argv)<3:
|
||
print("Too few arguments!", file=sys.stderr)
|
||
print(USAGE, file=sys.stderr)
|
||
sys.exit(USAGE_ERROR)
|
||
|
||
storage_directory=sys.argv[2]
|
||
if os.path.exists(storage_directory):
|
||
print("Target path already exists. Please choose a different location.", file=sys.stderr)
|
||
sys.exit(GENERAL_ERROR)
|
||
|
||
# default options
|
||
parity = False
|
||
parity_bytes = 512
|
||
checksum_algorithm = "sha512"
|
||
compress = False
|
||
# check for command line options
|
||
if len(sys.argv)>3:
|
||
arguments = sys.argv[3:]
|
||
for argument in arguments:
|
||
if not len(argument.split("="))==2:
|
||
print("Arguments to \"create\" always follow the scheme <option>=<value>.", file=sys.stderr)
|
||
print(USAGE, file=sys.stderr)
|
||
sys.exit(USAGE_ERROR)
|
||
option = argument.split("=")[0]
|
||
value = argument.split("=")[1]
|
||
if not option in ["parity", "parity-bytes", "checksum-algorithm", "compress"]:
|
||
print("Unknown option: "+option, file=sys.stderr)
|
||
print(USAGE, file=sys.stderr)
|
||
sys.exit(USAGE_ERROR)
|
||
if option=="parity":
|
||
if not value in ["on", "off"]:
|
||
print("Option \"parity\" accepts either \"on\" or \"off\".", file=sys.stderr)
|
||
sys.exit(USAGE_ERROR)
|
||
if value=="on":
|
||
parity = True
|
||
if option=="parity-bytes":
|
||
try:
|
||
parity_bytes = int(value)
|
||
except ValueError:
|
||
print("Option \"parity-bytes\" only accepts integers.", file=sys.stderr)
|
||
sys.exit(USAGE_ERROR)
|
||
if option=="checksum-algorithm":
|
||
if not value in hashlib.algorithms_available:
|
||
print("Chacksum algorithm \""+value+"\" not available.")
|
||
sys.exit(USAGE_ERROR)
|
||
checksum_algorithm = value
|
||
if option=="compress":
|
||
if not value in ["on", "off"]:
|
||
print("Option \"compress\" accepts either \"on\" or \"off\".", file=sys.stderr)
|
||
sys.exit(USAGE_ERROR)
|
||
if value=="on":
|
||
compress = True
|
||
|
||
create_container(storage_directory, parity=parity, parity_bytes=parity_bytes, checksum_algorithm=checksum_algorithm, compress=compress)
|
||
sys.exit(0)
|
||
|
||
# add subcommand: add a file to the storage container or add tags to it
|
||
# arguments:
|
||
# <storage directory> <hash|file> <tags ...>
|
||
# modifiers:
|
||
# file - requires a file path; adds a new file (or if file already in storage adds tags to that file), checks for collisions by comparing file size
|
||
# move - requires file modifier; moves the file to the storage dir instead of copying it
|
||
if command[0] == "add":
|
||
if len(sys.argv)<5:
|
||
print("Too few arguments!", file=sys.stderr)
|
||
print(USAGE, file=sys.stderr)
|
||
sys.exit(USAGE_ERROR)
|
||
|
||
storage_directory = sys.argv[2]
|
||
status, parity, parity_bytes, checksum_algorithm, compress = load_container_settings(storage_directory)
|
||
if not status==0:
|
||
if status==PATH_ERROR:
|
||
print("Invalid storage directory!", file=sys.stderr)
|
||
print(USAGE, file=sys.stderr)
|
||
if status==GENERAL_ERROR:
|
||
print("Verifying container settings failed.", file=sys.stderr)
|
||
sys.exit(status)
|
||
|
||
file_hash_or_path = sys.argv[3]
|
||
hash_allowed=True
|
||
if 'file' in command:
|
||
hash_allowed=False
|
||
if not any([hash_allowed and file_hash_or_path_is_known_hash(storage_directory, file_hash_or_path, compress), os.path.isfile(file_hash_or_path)]):
|
||
print("Unknown file!", file=sys.stderr)
|
||
print(USAGE, file=sys.stderr)
|
||
sys.exit(PATH_ERROR)
|
||
|
||
tags = sys.argv[4:]
|
||
|
||
if hash_allowed and file_hash_or_path_is_known_hash(storage_directory, file_hash_or_path, compress):
|
||
file_hash = file_hash_or_path
|
||
print("Hash for file in storage: "+file_hash)
|
||
else:
|
||
file_hash = calculate_file_hash(checksum_algorithm, file_hash_or_path)
|
||
if file_hash_or_path_is_known_hash(storage_directory, file_hash, compress):
|
||
print("File already in storage.")
|
||
#this assumes that the storage directory has not been tampered with or corrupted, FIXME!
|
||
if 'move' in command:
|
||
print("Removing external file.")
|
||
os.remove(file_hash_or_path)
|
||
else:
|
||
if 'move' in command:
|
||
print("Moving file to storage.")
|
||
shutil.move(file_hash_or_path, os.path.join(storage_directory, "objects", file_hash))
|
||
else:
|
||
print("Copying file to storage.")
|
||
shutil.copyfile(file_hash_or_path, os.path.join(storage_directory, "objects", file_hash))
|
||
if parity:
|
||
gegerate_parity_file(os.path.join(storage_directory, "objects", file_hash), parity_bytes, os.path.join(storage_directory, "parity", file_hash))
|
||
if compress:
|
||
print("Compressing...")
|
||
xz_process = os.popen("xz --best -T0 "+os.path.join(storage_directory, "objects", file_hash))
|
||
if parity:
|
||
xz_process = os.popen("xz --best -T0 "+os.path.join(storage_directory, "parity", file_hash))
|
||
|
||
for tag in tags:
|
||
print("Adding tag: "+tag)
|
||
if add_tag(storage_directory, file_hash, tag):
|
||
print("Added.")
|
||
else:
|
||
print("Tag already present.")
|
||
|
||
# lookup subcommand: return hash and tags of found files
|
||
# arguments: <storage directory> <hash|exact tag|set of exact tags>
|
||
# modifiers:
|
||
# first - only return one file
|
||
# unique - return error if not found or multiple found
|
||
# hash - perform lookup by hash
|
||
# tags - perform lookup by tag or set of tags
|
||
if command[0] == "lookup":
|
||
if len(sys.argv)<4:
|
||
print("Too few arguments!", file=sys.stderr)
|
||
print(USAGE, file=sys.stderr)
|
||
sys.exit(USAGE_ERROR)
|
||
|
||
storage_directory = sys.argv[2]
|
||
status, parity, parity_bytes, checksum_algorithm, compress = load_container_settings(storage_directory)
|
||
if not status==0:
|
||
if status==PATH_ERROR:
|
||
print("Invalid storage directory!", file=sys.stderr)
|
||
print(USAGE, file=sys.stderr)
|
||
if status==GENERAL_ERROR:
|
||
print("Verifying container settings failed.", file=sys.stderr)
|
||
sys.exit(status)
|
||
|
||
file_tags_or_hash = sys.argv[3:]
|
||
lookup_results = {}
|
||
if not 'tags' in command:
|
||
if file_hash_or_path_is_known_hash(storage_directory, file_tags_or_hash[0], compress):
|
||
lookup_results[file_tags_or_hash[0]] = get_tags_by_hash(storage_directory, file_tags_or_hash[0])
|
||
|
||
if not 'hash' in command:
|
||
# create a two dimensional array of all the files associated with each individual tag
|
||
file_hash_lists = []
|
||
for tag in file_tags_or_hash:
|
||
file_hash_lists = file_hash_lists + [get_hashes_by_tag(storage_directory, tag)]
|
||
# take the first of the arrays in the two dimensional array
|
||
common_file_hashes = file_hash_lists[0]
|
||
# iterate over the two dimensional array
|
||
for file_hash_list in file_hash_lists:
|
||
# check each element in common_file_hashes to ensure it is also in all other arrays in the two dimensional array, remove if it isn’t
|
||
for file_hash in common_file_hashes:
|
||
if not file_hash in file_hash_list:
|
||
common_file_hashes.remove(file_hash)
|
||
|
||
for file_hash in common_file_hashes:
|
||
lookup_results[file_hash] = get_tags_by_hash(storage_directory, file_hash)
|
||
|
||
if 'unique' in command:
|
||
if(len(lookup_results)==1):
|
||
print(json.dumps(lookup_results))
|
||
else:
|
||
print("More than one matching file found.", file=sys.stderr)
|
||
sys.exit(GENERAL_ERROR)
|
||
elif 'first' in command:
|
||
file_hash, tags = list(lookup_results.items())[0]
|
||
print(json.dumps({file_hash : tags}))
|
||
else:
|
||
print(json.dumps(lookup_results))
|
||
|
||
# link subcommand: add a symlink in <location> that points to the referenced file
|
||
# arguments:
|
||
# <storage directory> <hash> <location>
|
||
if command[0] == "link":
|
||
if len(sys.argv)<5:
|
||
print("Too few arguments!", file=sys.stderr)
|
||
print(USAGE, file=sys.stderr)
|
||
sys.exit(USAGE_ERROR)
|
||
|
||
storage_directory = sys.argv[2]
|
||
file_hash = sys.argv[3]
|
||
link_location = sys.argv[4]
|
||
status, parity, parity_bytes, checksum_algorithm, compress = load_container_settings(storage_directory)
|
||
if not status==0:
|
||
if status==PATH_ERROR:
|
||
print("Invalid storage directory!", file=sys.stderr)
|
||
print(USAGE, file=sys.stderr)
|
||
if status==GENERAL_ERROR:
|
||
print("Verifying container settings failed.", file=sys.stderr)
|
||
sys.exit(status)
|
||
|
||
if file_hash_or_path_is_known_hash(storage_directory, file_hash, compress):
|
||
if os.path.isdir(os.sep.join(link_location.split(os.sep)[:-1])):
|
||
if os.path.exists(link_location):
|
||
print(link_location+": file already exists.", file=sys.stderr)
|
||
sys.exit(GENERAL_ERROR)
|
||
else:
|
||
suffix = ""
|
||
if compress:
|
||
suffix = ".xz"
|
||
object_path = os.path.join(storage_directory, "objects", file_hash+suffix)
|
||
os.symlink(object_path, link_location)
|
||
print(link_location+" -> "+object_path)
|
||
else:
|
||
print("Parent directory "+os.sep.join(link_location.split(os.sep)[:-1])+" does not exist.", file=sys.stderr)
|
||
sys.exit(GENERAL_ERROR)
|
||
|
||
# help subcommand
|
||
if command[0] == "help":
|
||
print(USAGE)
|
||
sys.exit(0)
|
||
|
||
|
||
|
||
# this line is here to work around a bug in Xed
|