def _setup_ufw(self):
""" Enable UFW to controll the firewall
"""
print(Avalon.FM.BD + '\nEnable UFW firewall?' + Avalon.FM.RST)
print('Do you want SCUTUM to help configuring and enabling UFW firewall?')
print('This may help preventing a lot of scanning and attacks')
if Avalon.ask('Enable?', True):
# If ufw is not installed
if shutil.which('ufw') is None:
if Avalon.ask('UFW is not installed. Install?', True):
Utilities.install_packages(['ufw'])
else:
Avalon.warning('UFW package not available, disabling UFW')
self.config['Ufw']['handled'] = False
return
ufwctrl = Ufw()
print('SCUTUM can configure UFW Firewall for you')
print('However this will reset your current UFW configurations')
print('It is recommended to do so the first time you install SCUTUM')
if Avalon.ask('Let SCUTUM configure UFW for you?', True):
ufwctrl.initialize(True)
else:
Avalon.info('Okay. Then we will simply enable it for you')
ufwctrl.enable()
print('If you let SCUTUM handle UFW, then UFW will be activated and deactivated with SCUTUM')
if Avalon.ask('Let SCUTUM handle UFW?', True):
self.config['Ufw']['handled'] = True
else:
self.config['Ufw']['handled'] = False
else:
self.config['Ufw']['handled'] = False
Avalon.info('You can turn it on whenever you change your mind')
def check_memory():
""" Check usable system memory
Warn the user if insufficient memory is available for
the number of threads that the user have chosen.
"""
memory_status = []
# get system available memory
system_memory_available = psutil.virtual_memory().available / (1024 ** 3)
memory_status.append(('system', system_memory_available))
# check if Nvidia-smi is available
# GPUtil requires nvidia-smi.exe to interact with GPU
if args.method == 'gpu' or args.method == 'cudnn':
if not (shutil.which('nvidia-smi') or
os.path.isfile('C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe')):
# Nvidia System Management Interface not available
Avalon.warning('Nvidia-smi not available, skipping available memory check')
Avalon.warning('If you experience error \"cudaSuccess out of memory\", try reducing number of threads you\'re using')
else:
try:
# "0" is GPU ID. Both waifu2x drivers use the first GPU available, therefore only 0 makes sense
gpu_memory_available = (GPUtil.getGPUs()[0].memoryTotal - GPUtil.getGPUs()[0].memoryUsed) / 1024
memory_status.append(('GPU', gpu_memory_available))
except ValueError:
pass
# go though each checkable memory type and check availability
for memory_type, memory_available in memory_status:
if memory_type == 'system':
mem_per_thread = SYS_MEM_PER_THREAD
else:
mem_per_thread = GPU_MEM_PER_THREAD
# if user doesn't even have enough memory to run even one thread
if memory_available < mem_per_thread:
Avalon.warning(f'You might have insufficient amount of {memory_type} memory available to run this program ({memory_available} GB)')
Avalon.warning('Proceed with caution')
if args.threads > 1:
if Avalon.ask('Reduce number of threads to avoid crashing?', default=True, batch=args.batch):
args.threads = 1
# if memory available is less than needed, warn the user
elif memory_available < (mem_per_thread * args.threads):
Avalon.warning(f'Each waifu2x-caffe thread will require up to {SYS_MEM_PER_THREAD} GB of system memory')
Avalon.warning(f'You demanded {args.threads} threads to be created, but you only have {round(memory_available, 4)} GB {memory_type} memory available')
Avalon.warning(f'{mem_per_thread * args.threads} GB of {memory_type} memory is recommended for {args.threads} threads')
Avalon.warning(f'With your current amount of {memory_type} memory available, {int(memory_available // mem_per_thread)} threads is recommended')
# ask the user if he / she wants to change to the recommended
# number of threads
if Avalon.ask('Change to the recommended value?', default=True, batch=args.batch):
args.threads = int(memory_available // mem_per_thread)
else:
Avalon.warning('Proceed with caution')
def upgrade_full(self):
""" upgrade all packages
This method checks if there are packages
available for updating and update the packages
if updating them won't remove any packages from
the system. Often times when a bad source is added
to the system, APT tends to remove a number of packages
from the system when upgrading which is very risky.
"""
Avalon.info("Starting automatic upgrade")
Avalon.info("Updating APT cache")
with open("/etc/apt/sources.list", "r") as aptlist:
for line in aptlist:
if ("ubuntu.com" in line
and distro.linux_distribution()[0] != "Ubuntu"
and line.replace(" ", "")[0] != "#"):
Avalon.warning("Ubuntu source detected in source.list!")
Avalon.warning(
"Continue upgrading might cause severe consequences!")
if Avalon.ask("Are you sure that you want to continue?",
False):
break
else:
Avalon.warning("Aborting system upgrade..")
sys.exit(0)
self.update()
Avalon.info("APT cache updated")
if len(self.import_list) != 0:
Avalon.ask(
f"Detected un-imported keys: {' '.join(self.import_list)}")
# if there are no upgrades available
Avalon.info("Checking package updates")
if self.no_upgrades():
Avalon.debug_info("No upgrades available")
# if upgrades are available
else:
Avalon.info("Checking if full upgrade is safe")
# if upgrade is safe, use -y flag on apt-get full-upgrade
# otherwise, let user confirm the upgrade
if self.full_upgrade_safe():
Avalon.debug_info("Full upgrade is safe")
Avalon.info("Starting APT full upgrade")
self.full_upgrade()
else:
Avalon.warning("Full upgrade is NOT safe")
Avalon.warning("Requiring human confirmation")
self.manual_full_upgrade()
def enroll_settings():
settings = {}
settings['waifu2x_path'] = get_path('waifu2x-caffe-cui.exe path: ')
settings['ffmpeg_path'] = get_path('ffmpeg binaries directory: ')
settings['ffmpeg_arguments'] = []
while True:
argument = Avalon.gets(
'Extra arguments passed to ffmpeg (empty=none): ')
if argument:
settings['ffmpeg_arguments'].append(argument)
else:
break
settings['ffmpeg_hwaccel'] = Avalon.gets(
'ffmpeg hardware acceleration method (empty=auto): ')
if settings['ffmpeg_hwaccel'] == '':
settings['ffmpeg_hwaccel'] = 'auto'
settings['video2x_cache_folder'] = Avalon.gets(
'Video2X cache folder (empty=system default): ')
if settings['video2x_cache_folder'] == '':
settings['video2x_cache_folder'] = False
settings['preserve_frames'] = Avalon.ask(
'Preserve extracted or upscaled frames')
return settings
def enroll_settings():
settings = {}
settings['waifu2x_path'] = get_path('waifu2x-caffe-cui.exe path: ')
settings['ffmpeg_path'] = get_path('ffmpeg binaries directory: ')
settings['ffmpeg_arguments'] = []
while True:
argument = Avalon.gets('Extra arguments passed to ffmpeg (empty when done): ')
if argument:
settings['ffmpeg_arguments'].append(argument)
else:
break
settings['ffmpeg_hwaccel'] = Avalon.gets('ffmpeg hardware acceleration method (cuda): ')
if settings['ffmpeg_hwaccel'] == '':
settings['ffmpeg_hwaccel'] = 'cuda'
settings['extracted_frames'] = Avalon.gets('Temporary directory for extracted frames (empty for mkdtemp): ')
if settings['extracted_frames'] == '':
settings['extracted_frames'] = False
settings['upscaled_frames'] = Avalon.gets('Temporary directory for upscaled frames (empty for mkdtemp): ')
if settings['upscaled_frames'] == '':
settings['upscaled_frames'] = False
settings['preserve_frames'] = Avalon.ask('Preserve extracted or upscaled frames')
return settings
def json_save_profile(self, profile_path):
""" Save current profile to a JSON file
"""
# If profile already exists (file or link), ask the user if
# we should overwrite it.
if os.path.isfile(profile_path) or os.path.islink(profile_path):
if not Avalon.ask('File already exists. Overwrite?', True):
Avalon.warning('Aborted saving profile')
return 1
# Abort if profile_path points to a directory
if os.path.isdir(profile_path):
Avalon.warning('Destination path is a directory')
Avalon.warning('Aborted saving profile')
return 1
# Finally, write the profile into the destination file
Avalon.debug_info(f'Writing profile to: {profile_path}')
peers_dict = {}
peers_dict['peers'] = []
for peer in pm.peers:
peers_dict['peers'].append(peer.__dict__)
with open(profile_path, 'w') as profile:
json.dump(peers_dict, profile, indent=4)
profile.close()
def save_profile(self, profile_path):
""" Save current profile to a file
Serializes the current profile with pickle
and dumps it into a file.
"""
# If profile already exists (file or link), ask the user if
# we should overwrite it.
if os.path.isfile(profile_path) or os.path.islink(profile_path):
if not Avalon.ask('File already exists. Overwrite?', True):
Avalon.warning('Aborted saving profile')
return 1
# Abort if profile_path points to a directory
if os.path.isdir(profile_path):
Avalon.warning('Destination path is a directory')
Avalon.warning('Aborted saving profile')
return 1
# Finally, write the profile into the destination file
Avalon.debug_info('Writing profile to: {}'.format(profile_path))
with open(profile_path, 'wb') as profile:
pickle.dump(pm.peers, profile)
profile.close()
def uninstall(self):
""" Uninstall everything this system has deployed
This is yet to be completed.
"""
if not Avalon.ask(
'Are you sure to uninstall defense matrix completely?'):
return
# Restore original passwd binary
# passwd.restore_original_passwd()
# Remove SCUTUM
Utilities.execute(['scutum', '--uninstall'])
# Flush iptables and arptables
Utilities.execute(['iptables', '-F'])
Utilities.execute(['arptables', '-F'])
# Remove defense matrix directory
shutil.rmtree(self.install_dir)
# Remove linked executables
os.remove(self.executable)
exit(0)
def install_wicd_scripts(self):
""" Write scutum scripts for WICD
"""
print(Avalon.FG.G + '[+] INFO: Installing for WICD' + Avalon.FM.RST + '.....', end='')
if not os.path.isdir('/etc/wicd/'):
print(Avalon.FG.R + Avalon.FM.BD + 'ERROR' + Avalon.FM.RST)
Avalon.warning('WICD folder not found! WICD does not appear to be installed!')
if Avalon.ask('Continue anyway? (Create Directories)', False):
Utilities.execute(['mkdir', '-p', '/etc/wicd/scripts/postconnect/'])
Utilities.execute(['mkdir', '-p', '/etc/wicd/scripts/postdisconnect/'])
else:
Avalon.warning('Aborting installation for WICD')
return False
with open('/etc/wicd/scripts/postconnect/scutum_connect', 'w') as postconnect:
postconnect.write('#!/bin/bash\n')
postconnect.write('scutum')
postconnect.close()
with open('/etc/wicd/scripts/postdisconnect/scutum_disconnect', 'w') as postdisconnect:
postdisconnect.write('#!/bin/bash\n')
postdisconnect.write('scutum --reset')
postdisconnect.close()
Utilities.execute(['chown', 'root:', '/etc/wicd/scripts/postconnect/scutum_connect'])
Utilities.execute(['chmod', '755', '/etc/wicd/scripts/postconnect/scutum_connect'])
Utilities.execute(['chown', 'root:', '/etc/wicd/scripts/postdisconnect/scutum_disconnect'])
Utilities.execute(['chmod', '755', '/etc/wicd/scripts/postdisconnect/scutum_disconnect'])
print(Avalon.FG.G + Avalon.FM.BD + 'SUCCEED' + Avalon.FM.RST)
return True
def _install_rkhunter(self):
""" Install rkhunter
rkhunter is a rootkit detector that runs sanity checks
on system binary files and system misconfigurations.
"""
if not shutil.which('rkhunter'):
if Avalon.ask('rkhunter not installed. Install?', True):
self.pm_installation_list.append('rkhunter')
def _install_tigher(self):
""" Install tiger
Tiger is a package that will help controlling tripwire,
an HIDS, which hardens the system from the binary aspect.
"""
if not shutil.which('tiger'):
if Avalon.ask('tiger not installed. Install?', True):
self.pm_installation_list.append('tiger')
def _install_scutum_gui(self):
""" Install SCUTUM GUI
"""
print(Avalon.FM.BD + '\nInstall SCUTUM GUI?' + Avalon.FM.RST)
print('SCUTUM GUI is convenient for GUI Interfaces')
print('ex. KDE, GNOME, XFCE, etc.')
print('However, there\'s not point to install GUI on servers')
if Avalon.ask('Install SCUTUM GUI?', True):
if os.path.islink(self.DESKTOP_FILE) or os.path.isfile(self.DESKTOP_FILE):
os.remove(self.DESKTOP_FILE)
Utilities.execute(['ln', '-s', f'{self.INSTALL_DIR}/res/scutum-gui.desktop', self.DESKTOP_FILE])
def new_profile(self):
""" Create new profile and flush the peers list
"""
# Warn the user before flushing configurations
Avalon.warning('This will flush the currently loaded profile!')
if len(self.peers) != 0:
if not Avalon.ask('Continue?', False):
return
# Reset self.peers and start enrolling new peer data
self.peers = []
def add_peer():
""" Enroll a new peer
Gets all the information needed to generate a
new Peer class object.
"""
# Get peer tunnel address
while True:
address = Avalon.gets('Address (leave empty if client only): ')
result = re.match('^(?:\d{1,3}\.){3}\d{1,3}/{1}(?:\d\d?)?$', address)
if result is None:
Avalon.error('Invalid address entered')
Avalon.error('Please use CIDR notation (e.g. 10.0.0.0/8)')
continue
break
# Get peer public IP address
while True:
public_address = Avalon.gets(
'Public address (leave empty if client only): ')
result = re.match('^(?:\d{1,3}\.){3}\d{1,3}(?:/\d\d?)?$',
public_address)
if result is None and public_address != '': # field not required
Avalon.error('Invalid IP address entered')
continue
break
# Get peer listening port
listen_port = Avalon.gets('Listen port (leave empty for client): ')
# Get peer private key
private_key = Avalon.gets(
'Private key (leave empty for auto generation): ')
if private_key == '':
private_key = wg.genkey()
# Ask if this peer needs to be actively connected
# if peer is behind NAT and needs to be accessed actively
# PersistentKeepalive must be turned on (!= 0)
keep_alive = Avalon.ask('Keep alive?', False)
"""
preshared_key = False
if Avalon.ask('Use a preshared key?', True):
preshared_key = Avalon.gets('Preshared Key (leave empty for auto generation): ')
if preshared_key == '':
preshared_key = wg.genpsk()
peer = Peer(address, private_key, keep_alive, listen_port, preshared_key)
"""
peer = Peer(address, public_address, listen_port, private_key, keep_alive)
pm.peers.append(peer)
print_peer_config(peer)
def check_system_memory():
""" Check usable system memory
Warn the user if insufficient memory is available for
the number of threads that the user have chosen.
"""
memory_available = psutil.virtual_memory().available / (1024**3)
# If user doesn't even have enough memory to run even one thread
if memory_available < MEM_PER_THREAD:
Avalon.warning(
'You might have an insufficient amount of memory available to run this program ({} GB)'
.format(memory_available))
Avalon.warning('Proceed with caution')
if args.threads > 1:
if Avalon.ask('Reduce number of threads to avoid crashing?', True):
args.threads = 1
# If memory available is less than needed, warn the user
elif memory_available < (MEM_PER_THREAD * args.threads):
Avalon.warning(
'Each waifu2x-caffe thread will require up to 2.5 GB during initialization'
)
Avalon.warning(
'You demanded {} threads to be created, but you only have {} GB memory available'
.format(args.threads, round(memory_available, 4)))
Avalon.warning('{} GB of memory is recommended for {} threads'.format(
MEM_PER_THREAD * args.threads, args.threads))
Avalon.warning(
'With your current amount of memory available, {} threads is recommended'
.format(int(memory_available // MEM_PER_THREAD)))
# Ask the user if he / she wants to change to the recommended
# number of threads
if Avalon.ask('Change to the recommended value?', True):
args.threads = int(memory_available // MEM_PER_THREAD)
else:
Avalon.warning('Proceed with caution')
def _install_arp_controller_driver(self):
""" Install the CLI tool if not installed
"""
if self.config['ArpController']['driver'] == 'nftables':
binary = 'nft'
elif self.config['ArpController']['driver'] == 'arptables':
binary = 'arptables'
if shutil.which(binary) is None:
Avalon.warning('ARP controller driver is not installed')
if Avalon.ask(f'Install {self.config["ArpController"]["driver"]}?', True):
Utilities.install_packages([self.config['ArpController']['driver']])
else:
Avalon.error('ARP controller driver not installed')
Avalon.error('SCUTUM relies on the driver to run')
Avalon.error('Aborting installation')
exit(1)
def xinstall(packages: list):
""" install only packages that are not installed
By using the xinstall function to install packages,
already-installed packages will not get marked as
manually installed by APT.
Arguments:
packages {list} -- list of packages to install
"""
not_installed = []
installed = []
# get a list of all locally installed packages
apt_list = subprocess.run(
['apt', 'list'],
stderr=subprocess.DEVNULL).stdout.decode().split('\n')
for line in apt_list:
for package in packages:
if package == line.split('/')[0] and 'installed' not in line:
not_installed.append(package)
elif package == line.split('/')[0] and 'installed' in line:
installed.append(package)
Avalon.info('Packages already installed:')
print(' '.join(installed))
Avalon.info('Packages to be installed:')
print(' '.join(not_installed))
execute = ['apt-get', 'install', '-s']
execute.extend(not_installed)
Avalon.info('Launching a dry-run')
subprocess.call(execute)
if Avalon.ask('Confirm installation:', False):
# swap -s flag with -y for actual installation
execute[execute.index('-s')] = '-y'
subprocess.call(execute)
else:
Avalon.warning('Installation aborted')
def command_interpreter(db_connection, commands):
""" AnyRadius shell command interpreter
"""
try:
# Try to guess what the user is saying
possibilities = [
s for s in COMMANDS if s.lower().startswith(commands[1])
]
if len(possibilities) == 1:
commands[1] = possibilities[0]
if commands[1].replace(' ', '') == '':
result = 0
elif commands[1].lower() == 'help':
print_help()
result = 0
elif commands[1].lower() == 'truncateusertable':
Avalon.warning('By truncating you will LOSE ALL USER DATA')
if Avalon.ask('Are you sure you want to truncate?'):
result = db_connection.truncate_user_table()
else:
Avalon.warning('Operation canceled')
result = 0
elif commands[1].lower() == 'adduser':
result = db_connection.add_user(commands[2], commands[3])
elif commands[1].lower() == 'deluser':
result = db_connection.del_user(commands[2])
elif commands[1].lower() == 'showusers':
result = db_connection.show_users()
elif commands[1].lower() == 'exit' or commands[1].lower() == 'quit':
Avalon.warning('Exiting')
exit(0)
elif len(possibilities) > 0:
Avalon.warning('Ambiguous command \"{}\"'.format(commands[1]))
print('Use \"Help\" command to list available commands')
result = 1
else:
Avalon.error('Invalid command')
print('Use \"Help\" command to list available commands')
result = 1
return result
except IndexError:
Avalon.error('Invalid arguments')
print('Use \"Help\" command to list available commands')
result = 0
def install_nm_scripts(self, interfaces):
""" Write scutum scripts for Network Manager
"""
print(Avalon.FG.G + '[+] INFO: Installing for NetworkManager' + Avalon.FM.RST + '.....', end='')
if not os.path.isdir('/etc/NetworkManager/dispatcher.d/'):
print(Avalon.FG.R + Avalon.FM.BD + 'ERROR' + Avalon.FM.RST)
Avalon.warning('NetworkManager folders not found! NetworkManager does not appear to be installed!')
if Avalon.ask('Continue anyway? (Create Directories)', False):
Utilities.execute(['mkdir', '-p', '/etc/NetworkManager/dispatcher.d/'])
else:
Avalon.warning('Aborting installation for NetworkManager')
return False
with open('/etc/NetworkManager/dispatcher.d/scutum', 'w') as nmScript:
nmScript.write('#!/usr/bin/env python3\n')
nmScript.write('\n')
nmScript.write('import sys\n')
nmScript.write('import os\n')
nmScript.write('\n')
nmScript.write('try:\n')
nmScript.write(' interface = sys.argv[1]\n')
nmScript.write(' status = sys.argv[2]\n')
nmScript.write('except IndexError:\n')
nmScript.write(' exit(0)\n')
nmScript.write('\n')
nmScript.write('if status == \"down\":\n')
nmScript.write(' os.system(\"scutum --reset\")\n')
nmScript.write(' exit(0)\n')
nmScript.write('\n')
nmScript.write('if status == \"up\":\n')
nmScript.write(' os.system(\"scutum\")\n')
nmScript.write(' exit(0)\n')
nmScript.close()
Utilities.execute(['chown', 'root:', '/etc/NetworkManager/dispatcher.d/scutum'])
Utilities.execute(['chmod', '755', '/etc/NetworkManager/dispatcher.d/scutum'])
print(Avalon.FG.G + Avalon.FM.BD + 'SUCCEED' + Avalon.FM.RST)
return True
def _install_defense_matrix(self):
""" Installs defense matrix to system and
link defense matrix executable to bin path.
"""
# Get installation destination from user
user_install_dir = Avalon.gets(
'Installation destination (\"/usr/share/defense-matrix\")')
if user_install_dir != '':
self.install_dir = user_install_dir
# If files already at the correct directory, pass
if self.current_dir == self.install_dir:
pass
# Check if destination directory occupied
else:
if os.path.isdir(self.install_dir) or os.path.islink(
self.install_dir):
if not Avalon.ask('Target directory exists. Overwrite?', True):
Avalon.warning(
'Aborting installation: target directory not writable')
if os.path.isdir(self.install_dir):
shutil.rmtree(self.install_dir)
else:
os.remove(self.install_dir)
# Copy defense matrix to destination directory
shutil.copytree(self.current_dir, self.install_dir)
# If defense-matrix is already linked to path, remove it
if os.path.islink(self.executable) or os.path.isfile(self.executable):
os.remove(self.executable) # Remove old file or symbolic links
# Link current defense-matrix.py to path
os.symlink('{}/bin/defense-matrix.py'.format(self.current_dir),
self.executable)
def check_version():
""" check if KPM is up-to-date
Check if KPM is up to date with the the newest
version on GitHub. Prompt the user to upgrade if
the local version is not the newest.
"""
# get version number of KPM on GitHub
Avalon.info("Checking KPM's Version")
latest = requests.get(
"https://api.github.com/repos/k4yt3x/kpm/releases/latest")
latest_json = latest.json()
# if rate limit is exceeded, 403 will be returned
if (latest.status_code == requests.codes.forbidden
and "API rate limit exceeded" in latest_json["message"]):
Avalon.warning("GitHub API rate limit exceeded")
return
# if the status code isn't 200, warn the user and skip the rest of the checks
elif latest.status_code != requests.codes.ok:
Avalon.warning("GitHub API request encountered an error")
return
latest_version = latest_json["tag_name"]
Avalon.debug_info(f"Server version: {latest_version}")
# if the server version is newer than local version
if version.parse(latest_version) > version.parse(VERSION):
Avalon.info(
f"There is a newer version of KPM available ({latest_version})")
if Avalon.ask("Update to the newest version?", True):
upgrade_kpm()
else:
Avalon.warning("Skipping the upgrade for now")
else:
Avalon.debug_info(f"KPM is already on the newest version ({VERSION})")
def check_version():
""" check if KPM is up-to-date
Check if KPM is up to date with the the newest
version on GitHub. Prompt the user to upgrade if
the local version is not the newest.
"""
# get version number of KPM on GitHub
Avalon.debug_info('Checking KPM Version')
for line in requests.get(GITHUB_KPM_FILE).text.split('\n'):
if 'VERSION = ' in line:
server_version = line.split(' ')[-1].replace('\'', '')
break
Avalon.debug_info(f'Server version: {server_version}')
# if the server version is newer than local version
if server_version > VERSION:
Avalon.info('Here\'s a newer version of KPM!')
if Avalon.ask('Update to the newest version?', True):
upgrade_kpm()
else:
Avalon.warning('Ignoring update')
else:
Avalon.debug_info('KPM is already on the newest version')
raise ArgumentError('both scaling ration and width/height specified')
if (args.width and not args.height) or (not args.width and args.height):
Avalon.error('You must specify both width and height')
raise ArgumentError('only one of width or height is specified')
# check available memory if driver is waifu2x-based
if args.driver in [
'waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan'
]:
check_memory()
# anime4k runs significantly faster with more processes
if args.driver == 'anime4k' and args.processes <= 1:
Avalon.warning('Anime4K runs significantly faster with more processes')
if Avalon.ask('Use more processes of Anime4K?',
default=True,
batch=args.batch):
while True:
try:
processes = Avalon.gets('Amount of processes to use [5]: ',
default=5,
batch=args.batch)
args.processes = int(processes)
break
except ValueError:
if processes == '':
args.processes = 5
break
Avalon.error(f'{processes} is not a valid integer')
# read configurations from configuration file
raise ArgumentError('scaling ratio must be 1 or 2 for waifu2x_ncnn_vulkan')
if (args.width or args.height) and args.ratio:
Avalon.error('You can only specify either scaling ratio or output width and height')
raise ArgumentError('both scaling ration and width/height specified')
if (args.width and not args.height) or (not args.width and args.height):
Avalon.error('You must specify both width and height')
raise ArgumentError('only one of width or height is specified')
# check available memory if driver is waifu2x-based
if args.driver in ['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan']:
check_memory()
# anime4k runs significantly faster with more threads
if args.driver == 'anime4k' and args.threads <= 1:
Avalon.warning('Anime4K runs significantly faster with more threads')
if Avalon.ask('Use more threads of Anime4K?', True):
while True:
try:
threads = Avalon.gets('Amount of threads to use [5]: ')
args.threads = int(threads)
break
except ValueError:
if threads == '':
args.threads = 5
break
else:
Avalon.error(f'{threads} is not a valid integer')
# read configurations from JSON
config = read_config(args.config)
config = absolutify_paths(config)
def add_peer():
""" Enroll a new peer
Gets all the information needed to generate a
new Peer class object.
"""
# Get peer tunnel address
while True:
address = Avalon.gets(
'Address (leave empty if client only) [IP/CIDR]: ')
if re.match('^(?:\d{1,3}\.){3}\d{1,3}/{1}(?:\d\d?)?$',
address) is None:
Avalon.error('Invalid address entered')
Avalon.error('Please use CIDR notation (e.g. 10.0.0.0/8)')
continue
break
# Get peer public IP address
while True:
public_address = Avalon.gets(
'Public address (leave empty if client only) [IP|FQDN]: ')
# Check if public_address is valid IP or FQDN
valid_address = False
if re.match('^(?:\d{1,3}\.){3}\d{1,3}(?:/\d\d?)?$',
public_address) is not None:
valid_address = True
if re.match(
'(?=^.{4,253}$)(^((?!-)[a-zA-Z0-9-]{1,63}(?<!-)\.)+[a-zA-Z]{2,63}$)',
public_address) is not None:
valid_address = True
if not valid_address and public_address != '': # field not required
Avalon.error('Invalid public address address entered')
Avalon.error('Please enter an IP address or FQDN')
continue
break
# Get peer listening port
listen_port = Avalon.gets(
'Listen port (leave empty for client) [1-65535]: ')
# Get peer private key
private_key = Avalon.gets(
'Private key (leave empty for auto generation): ')
if private_key == '':
private_key = wg.genkey()
# Ask if this peer needs to be actively connected
# if peer is behind NAT and needs to be accessed actively
# PersistentKeepalive must be turned on (!= 0)
keep_alive = Avalon.ask('Keep alive?', False)
"""
preshared_key = False
if Avalon.ask('Use a preshared key?', True):
preshared_key = Avalon.gets('Preshared Key (leave empty for auto generation): ')
if preshared_key == '':
preshared_key = wg.genpsk()
peer = Peer(address, private_key, keep_alive, listen_port, preshared_key)
"""
# Get peer alias
alias = Avalon.gets('Alias (optional): ')
# Get peer description
description = Avalon.gets('Description (optional): ')
# Create peer and append peer into the peers list
peer = Peer(address,
public_address,
listen_port,
private_key,
keep_alive=keep_alive,
alias=alias,
description=description)
pm.peers.append(peer)
print_peer_config(peer)
def upgrade_all(self):
""" upgrade all packages
This method checks if there are packages
available for updating and update the packages
if updating them won't remove any packages from
the system. Often times when a bad source is added
to the system, APT tends to remove a number of packages
from the system when upgrading which is very risky.
"""
Avalon.info('Starting automatic upgrade')
Avalon.info('Updating APT cache')
with open('/etc/apt/sources.list', 'r') as aptlist:
for line in aptlist:
if 'ubuntu.com' in line and distro.linux_distribution(
)[0] != 'Ubuntu' and line.replace(' ', '')[0] != '#':
Avalon.warning('Ubuntu source detected in source.list!')
Avalon.warning(
'Continue upgrading might cause severe consequences!')
if Avalon.ask('Are you sure that you want to continue?',
False):
break
else:
Avalon.warning('Aborting system upgrade..')
exit(0)
self.update()
Avalon.info('APT cache updated')
if len(self.import_list) != 0:
if Avalon.ask('Detected unimported keys, import?', True):
if shutil.which('dirmngr') is None:
Avalon.warning('dirmngr Not installed')
Avalon.warning('It is required for importing keys')
# ask if user wants to install dirmngr
if Avalon.ask('Install Now?'):
self.install('dirnmgr')
# check dirmngr after package installation
if isinstance(shutil.which('dirmngr'), str):
Avalon.info('Installation successful')
self.import_keys(self.import_list)
Avalon.info('Keys imported')
Avalon.info(
'Updating APT cache after key importing')
self.update()
else:
Avalon.error('Installation Failed')
Avalon.error('Please check your settings')
Avalon.warning(
'dirmngr not available. Continuing without importing keys'
)
else:
Avalon.warning('dirmngr not available')
Avalon.warning('Continuing without importing keys')
else:
self.import_keys(self.import_list)
Avalon.info('Keys imported')
Avalon.info('Updating APT cache after key importing')
self.update()
# Second update after keys are imported
self.update()
# if there are no upgrades available
Avalon.debug_info('Checking package updates')
if self.no_upgrades():
Avalon.info('No upgrades available')
# if upgrades are available
else:
Avalon.debug_info('Checking if full upgrade is safe')
# if upgrade is safe, use -y flag on apt-get full-upgrade
# otherwise, let user confirm the upgrade
if self.full_upgrade_safe():
Avalon.info('Full upgrade is safe')
Avalon.info('Starting APT full upgrade')
self.full_upgrade()
else:
Avalon.warning('Full upgrade is NOT safe')
Avalon.warning('Requiring human confirmation')
self.manual_full_upgrade()
check_version()
# if -x, --xinstall specified
if args.xinstall:
packages = args.xinstall.split(',')
kobj.xinstall(packages)
# if no arguments are given
else:
kobj.upgrade_all()
Avalon.debug_info('Checking for unused packages')
# check if there are any unused packages
if kobj.autoremove_available():
if Avalon.ask('Remove useless packages?', True):
kobj.autoremove()
else:
Avalon.info('No unused packages found')
# apt autoclean
Avalon.info('Erasing old downloaded archive files')
kobj.autoclean()
except KeyboardInterrupt:
Avalon.warning('Aborting')
except Exception:
Avalon.error('Error caught during execution')
traceback.print_exc()
def generate_configs(output_path):
""" Generate configuration file for every peer
This function reads the PEERS list, generates a
configuration file for every peer, and export into
the CONFIG_OUTPUT directory.
"""
if len(pm.peers) == 0:
Avalon.warning('No peers configured, exiting')
exit(0)
if len(pm.peers) == 1:
Avalon.warning('Only one peer configured')
Avalon.info('Generating configuration files')
# Abort is destination is a file / link
if os.path.isfile(output_path) or os.path.islink(output_path):
Avalon.warning('Destination path is a file / link')
Avalon.warning('Aborting configuration generation')
return 1
# Ask if user wants to create the output directory if it doesn't exist
if not os.path.isdir(output_path):
if Avalon.ask(
'Output directory doesn\'t exist. Create output directory?',
True):
os.mkdir(output_path)
else:
Avalon.warning('Aborting configuration generation')
return 1
# Iterate through all peers and generate configuration for each peer
for peer in pm.peers:
Avalon.debug_info('Generating configuration file for {}'.format(
peer.address))
with open('{}/{}.conf'.format(output_path,
peer.address.split('/')[0]),
'w') as config:
# Write Interface configuration
config.write('[Interface]\n')
config.write('PrivateKey = {}\n'.format(peer.private_key))
if peer.address != '':
config.write('Address = {}\n'.format(peer.address))
if peer.listen_port != '':
config.write('ListenPort = {}\n'.format(peer.listen_port))
# Write peers' information
for p in pm.peers:
if p.address == peer.address:
# Skip if peer is self
continue
config.write('\n[Peer]\n')
print(p.private_key)
config.write('PublicKey = {}\n'.format(wg.pubkey(
p.private_key)))
config.write('AllowedIPs = {}\n'.format(p.address))
if p.public_address != '':
config.write('Endpoint = {}:{}\n'.format(
p.public_address, p.listen_port))
if peer.keep_alive:
config.write('PersistentKeepalive = 25\n')
if p.preshared_key:
config.write('PresharedKey = {}\n'.format(p.preshared_key))
# If memory available is less than needed, warn the user
elif memory_available < (MEM_PER_THREAD * args.threads):
Avalon.warning(
'Each waifu2x-caffe thread will require up to 2.5 GB during initialization'
)
Avalon.warning(
'You demanded {} threads to be created, but you only have {} GB memory available'
.format(args.threads, round(memory_available, 4)))
Avalon.warning('{} GB of memory is recommended for {} threads'.format(
MEM_PER_THREAD * args.threads, args.threads))
Avalon.warning(
'With your current amount of memory available, {} threads is recommended'
.format(int(memory_available // MEM_PER_THREAD)))
# Ask the user if he / she wants to change to the recommended
# number of threads
if Avalon.ask('Change to the recommended value?', True):
args.threads = int(memory_available // MEM_PER_THREAD)
else:
Avalon.warning('Proceed with caution')
# Start execution
try:
begin_time = time.time()
video2x()
Avalon.info('Program completed, taking {} seconds'.format(
round((time.time() - begin_time), 5)))
except Exception:
Avalon.error('An exception occurred')
traceback.print_exc()
# load video2x settings
video2x_cache_folder = config['video2x']['video2x_cache_folder']
preserve_frames = config['video2x']['preserve_frames']
# create temp directories if they don't exist
if not video2x_cache_folder:
video2x_cache_folder = '{}\\video2x'.format(tempfile.gettempdir())
if video2x_cache_folder and not os.path.isdir(video2x_cache_folder):
if not os.path.isfile(video2x_cache_folder) and not os.path.islink(
video2x_cache_folder):
Avalon.warning(
'Specified cache folder/directory {} does not exist'.format(
video2x_cache_folder))
if Avalon.ask('Create folder/directory?',
default=True,
batch=args.batch):
if os.mkdir(video2x_cache_folder) is None:
Avalon.info('{} created'.format(video2x_cache_folder))
else:
Avalon.error(
'Unable to create {}'.format(video2x_cache_folder))
Avalon.error('Aborting...')
exit(1)
else:
Avalon.error('Specified cache folder/directory is a file/link')
Avalon.error('Unable to continue, exiting...')
exit(1)
# start execution
try:
