monkey/infection_monkey/monkeyfs.py

62 lines
1.4 KiB
Python

from io import BytesIO
import os
__author__ = 'hoffer'
MONKEYFS_PREFIX = 'monkeyfs://'
open_orig = open
class VirtualFile(BytesIO):
_vfs = {} # virtual File-System
def __init__(self, name, mode='r', buffering=None):
if not name.startswith(MONKEYFS_PREFIX):
name = MONKEYFS_PREFIX + name
self.name = name
self._mode = mode
if name in VirtualFile._vfs:
super(VirtualFile, self).__init__(self._vfs[name])
else:
super(VirtualFile, self).__init__('')
def flush(self):
super(VirtualFile, self).flush()
VirtualFile._vfs[self.name] = self.getvalue()
@staticmethod
def getsize(path):
return len(VirtualFile._vfs[path])
@staticmethod
def isfile(path):
return path in VirtualFile._vfs
def getsize(path):
if path.startswith(MONKEYFS_PREFIX):
return VirtualFile.getsize(path)
else:
return os.stat(path).st_size
def isfile(path):
if path.startswith(MONKEYFS_PREFIX):
return VirtualFile.isfile(path)
else:
return os.path.isfile(path)
def virtual_path(name):
return "%s%s" % (MONKEYFS_PREFIX, name)
def open(name, mode='r', buffering=-1):
# use normal open for regular paths, and our "virtual" open for monkeyfs:// paths
if name.startswith(MONKEYFS_PREFIX):
return VirtualFile(name, mode, buffering)
else:
return open_orig(name, mode=mode, buffering=buffering)