Added/changed docstrings to FileBase

pull/7/head
Dan Puttick 2016-12-05 21:02:46 -05:00
parent 146f258e54
commit 62a7f680b4
1 changed files with 45 additions and 32 deletions

View File

@ -1,5 +1,13 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Contains the base objects for use when creating a sanitizer using
PyCIRCLean. Subclass FileBase and KittenGroomerBase to implement your
desired behavior.
"""
import os
import sys
import hashlib
@ -11,36 +19,36 @@ from twiggy import quick_setup, log
class KittenGroomerError(Exception):
"""Base KittenGroomer exception handler."""
def __init__(self, message):
'''
Base KittenGroomer exception handler.
'''
super(KittenGroomerError, self).__init__(message)
self.message = message
class ImplementationRequired(KittenGroomerError):
'''
Implementation required error
'''
"""Implementation required error."""
pass
class FileBase(object):
"""
Base object for individual files in the source directory. Has information
about the file as attributes and various helper methods. Initialised with
the source path and expected destination path. Subclass and add attributes
or methods relevant to a given implementation."
"""
def __init__(self, src_path, dst_path):
'''
Contains base information for a file on the source USB key,
initialised with expected src and dest path
'''
self.src_path = src_path
self.dst_path = dst_path
self.log_details = {'filepath': self.src_path}
self.log_string = ''
a, self.extension = os.path.splitext(self.src_path)
self._get_mimetype()
_, self.extension = os.path.splitext(self.src_path)
self._determine_mimetype()
def _get_mimetype(self):
def _determine_mimetype(self):
if os.path.islink(self.src_path):
# magic will throw an IOError on a broken symlink
self.mimetype = 'inode/symlink'
@ -55,7 +63,6 @@ class FileBase(object):
self.mimetype = mt.decode("utf-8")
except:
self.mimetype = mt
if self.mimetype and '/' in self.mimetype:
self.main_type, self.sub_type = self.mimetype.split('/')
else:
@ -63,40 +70,53 @@ class FileBase(object):
self.sub_type = ''
def has_mimetype(self):
"""
Returns True if file has a full mimetype, else False.
Returns False + updates log if self.main_type or self.sub_type
are not set.
"""
if not self.main_type or not self.sub_type:
self.log_details.update({'broken_mime': True})
return False
return True
def has_extension(self):
"""
Returns True if self.extension is set, else False.
Returns False + updates self.log_details if self.extension is not set.
"""
if not self.extension:
self.log_details.update({'no_extension': True})
return False
return True
def is_dangerous(self):
"""Returns True if self.log_details contains 'dangerous'."""
if self.log_details.get('dangerous'):
return True
return False
def is_symlink(self):
"""Returns True and updates log if file is a symlink."""
if self.has_mimetype() and self.main_type == 'inode' and self.sub_type == 'symlink':
self.log_details.update({'symlink': os.readlink(self.src_path)})
return True
return False
def add_log_details(self, key, value):
'''
Add an entry in the log dictionary
'''
"""Takes a key + a value and adds them to self.log_details."""
self.log_details[key] = value
def make_dangerous(self):
'''
This file should be considered as dangerous and never run.
Prepending and appending DANGEROUS to the destination
file name avoid double-click of death
'''
"""
Marks a file as dangerous.
Prepends and appends DANGEROUS to the destination file name
to avoid double-click of death.
"""
if self.is_dangerous():
# Already marked as dangerous, do nothing
return
@ -105,11 +125,7 @@ class FileBase(object):
self.dst_path = os.path.join(path, 'DANGEROUS_{}_DANGEROUS'.format(filename))
def make_unknown(self):
'''
This file has an unknown type and it was not possible to take
a decision. Theuser will have to decide what to do.
Prepending UNKNOWN
'''
"""Marks a file as an unknown type and prepends UNKNOWN to filename."""
if self.is_dangerous() or self.log_details.get('binary'):
# Already marked as dangerous or binary, do nothing
return
@ -118,11 +134,7 @@ class FileBase(object):
self.dst_path = os.path.join(path, 'UNKNOWN_{}'.format(filename))
def make_binary(self):
'''
This file is a binary, and should probably not be run.
Appending .bin avoir double click of death but the user
will have to decide by itself.
'''
"""Marks a file as a binary and appends .bin to filename."""
if self.is_dangerous():
# Already marked as dangerous, do nothing
return
@ -131,6 +143,7 @@ class FileBase(object):
self.dst_path = os.path.join(path, '{}.bin'.format(filename))
def force_ext(self, ext):
"""If dst_path does not end in ext, appends the ext and updates log."""
if not self.dst_path.endswith(ext):
self.log_details['force_ext'] = True
self.dst_path += ext