Add support for random configuration values

This commit is contained in:
Oleksii Shevchuk 2017-03-11 08:18:47 +02:00
parent e74aad2adc
commit 0497299528
1 changed files with 29 additions and 0 deletions

View File

@ -7,6 +7,8 @@ except ImportError:
from os import path, makedirs from os import path, makedirs
from netaddr import IPAddress from netaddr import IPAddress
import platform import platform
import random
import string
class PupyConfig(ConfigParser): class PupyConfig(ConfigParser):
def __init__(self, config='pupy.conf'): def __init__(self, config='pupy.conf'):
@ -21,6 +23,7 @@ class PupyConfig(ConfigParser):
self.project_path, self.project_path,
config config
] ]
self.randoms = {}
ConfigParser.__init__(self) ConfigParser.__init__(self)
self.read(self.files) self.read(self.files)
@ -72,8 +75,23 @@ class PupyConfig(ConfigParser):
else: else:
return path.abspath(retfolder) return path.abspath(retfolder)
def set(self, section, key, value):
if section != 'randoms':
ConfigParser.set(self, section, key, value)
else:
self.randoms[key] = value
def get(self, *args, **kwargs): def get(self, *args, **kwargs):
try: try:
if args[0] == 'randoms':
if not args[1] in self.randoms:
N = kwargs.get('random', 10)
self.randoms[args[1]] = ''.join(
random.choice(
string.ascii_letters + string.digits) for _ in range(N))
return self.randoms[args[1]]
return ConfigParser.get(self, *args, **kwargs) return ConfigParser.get(self, *args, **kwargs)
except Error as e: except Error as e:
return None return None
@ -83,3 +101,14 @@ class PupyConfig(ConfigParser):
if not ip: if not ip:
return None return None
return IPAddress(ip) return IPAddress(ip)
def sections(self):
sections = ConfigParser.sections(self)
sections.append('randoms')
return sections
def options(self, section):
if section != 'randoms':
return ConfigParser.options(self, section)
return self.randoms.keys()