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: