def uninstall(self): """ Remove SCUTUM completely This method will remove all SCUTUM files from the system. It is yet to be tested completely. """ self.remove_wicd_scripts() self.remove_nm_scripts() Utilities.execute(['ufw', '--force', 'reset']) # Reset ufw configurations Utilities.execute(['rm', '-f', '/etc/ufw/*.*.*']) # Delete automatic backups # A list of files, directories and links to remove rmlist = ['/usr/bin/scutum', self.INSTALL_DIR, self.CONFPATH, '/var/log/scutum.log', '/usr/lib/systemd/system/scutum.service', '/etc/init.d/scutum', '/etc/systemd/system/multi-user.target.wants/scutum.service' ] # Remove all files in rmlist for path in rmlist: if os.path.isfile(path) or os.path.islink(path): os.remove(path) elif os.path.isdir(path): shutil.rmtree(path) Avalon.info('SCUTUM removed successfully!') exit(0)
def install(self): """ Start the installation Start the installation and install all packages and components. """ # A list of packages to be installed # by system package manager self.pm_installation_list = [] self._install_defense_matrix() # self._install_passwdcmplx() # Packages to be installed by pm self._install_tigher() self._install_rkhunter() # Commit installation if len(self.pm_installation_list) > 0: Utilities.install_packages(self.pm_installation_list) # Install SCUTUM separately since it's not a # standard package self._install_scutum() print('\n' + Avalon.FM.BD, end='') Avalon.info('Defense Matrix installation completed') Avalon.info( 'You can now control it via the \"defense-matrix\" command')
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 install(self): """ Install SCUTUM This method will install all SCUTUM files and components """ Avalon.info('Starting installation procedure') # Initialize configuration containers self.config = {} self.config['Interfaces'] = {} self.config['NetworkControllers'] = {} self.config['Ufw'] = {} self.config['ArpController'] = {} self._install_scutum_files() self._install_service() self._get_arp_controller_driver() self._install_arp_controller_driver() self._get_controlled_interfaces() self._get_controlled_nm() self._setup_ufw() self._install_scutum_gui() # Export the configuration into configuration file with open(self.CONFPATH, 'w') as configfile: json.dump(self.config, configfile, indent=2) configfile.close() print('\n' + Avalon.FM.BD, end='') Avalon.info('SCUTUM installation completed') Avalon.info('SCUTUM service is now enabled on system startup') Avalon.info('You can now control it with systemd') Avalon.info('You can also control it manually via the \"scutum\" command')
def allow(self, port): """ Accept all traffic from one address Arguments: port {int} -- Port number """ Avalon.info(f'[UFW]: Allowing port {str(port)}\n\n') Utilities.execute(['ufw', 'allow', f'{str(port)}/tcp'])
def expire(self, port): """ Disallows all traffic from one address Arguments: port {int} -- Port number """ Avalon.info(f'[UFW]: Expiring port {str(port)}\n\n') Utilities.execute(['ufw', '--force', 'delete', 'allow', f'{str(port)}/tcp'])
def resolve(self): Avalon.info(f"正在对 {self._domain} 进行全球解析……") self._session.cookies.clear() self._get_src() self._get_token() self._get_dns_id() self._global_query() self._extend_query() Avalon.info(f"{self._domain} 的全球解析已完成")
def main(): """ AnyRadius Manager main function This function can only be executed when this file is not being imported. """ # Create database controller connection try: if sys.argv[1].lower() == 'help': print_help() exit(0) elif sys.argv[1].lower() == 'config': config_path = sys.argv[2] else: config_path = '/etc/anyradius.json' except IndexError: Avalon.error('Error parsing configuration file path') exit(1) Avalon.debug_info('Reading config from: {}'.format(config_path)) db_host, db_user, db_pass, db, table = read_config(config_path) Avalon.info('Connecting to RADIUS database') rdb = UserDatabase(db_host, db_user, db_pass, db, table) Avalon.info('Database connection established') # Begin command interpreting try: if sys.argv[1].lower() == 'interactive' or sys.argv[1].lower( ) == 'int': print_legal_info() # Set command completer completer = ShellCompleter(COMMANDS) readline.set_completer(completer.complete) readline.parse_and_bind('tab: complete') # Launch interactive trojan shell prompt = '{}[AnyRadius]> {}'.format(Avalon.FM.BD, Avalon.FM.RST) while True: command_interpreter(rdb, [''] + input(prompt).split(' ')) else: # Return to shell with command return value exit(command_interpreter(rdb, sys.argv[0:])) except IndexError: Avalon.warning('No commands specified') exit(0) except (MySQLdb.Error, MySQLdb.Warning): Avalon.error('Database error') traceback.print_exc() exit(1) except (KeyboardInterrupt, EOFError): Avalon.warning('Exiting') exit(0) except Exception: Avalon.error('Exception caught') traceback.print_exc() exit(1)
def show_users(self): """ List all users from the database """ total_users = self.cursor.execute("SELECT * FROM {}".format( self.table)) table = PrettyTable(['Username', 'Password']) for user in self.cursor.fetchall(): table.add_row([user[0], user[1], user[2], user[4]]) print(table) Avalon.info( 'Query complete, {} users found in database'.format(total_users))
def get_ip_ping_delay(ip): ping_result = ping(ip, count=5) average_delay = ping_result.rtt_avg_ms if average_delay < 100: Avalon.info(f"{ip}\t平均延迟: {average_delay} ms") else: Avalon.error(f"{ip}\t平均延迟: {average_delay} ms") return average_delay
def start(self): """ Dev: K4YT3X IZAYOI Date Created: Jan 15, 2018 Last Modified: Jan 16, 2018 Dev: Reimannsum Last Modified: Aug 27, 2019 This method is the main ISS pointer controller it runs infinitively until Ctrl^C is pressed. """ ts = load.timescale() stations_url = 'http://celestrak.com/NORAD/elements/stations.txt' satellites = load.tle(stations_url) satellite = satellites['ISS (ZARYA)'] observer = Topos('42.5337N', '83.7384W') while True: t = ts.now() days = t - satellite.epoch if abs(days) > 14: satellites = load.tle(stations_url, reload=True) satellite = satellites['ISS (ZARYA)'] self.pointer.check_gravity() difference = satellite - observer topocentric = difference.at(t) alt, az, distance = topocentric.altaz() self.pointer.elevation_set(alt.degrees) self.pointer.azimuth_set(az.degrees) self.display.set_pointing(az.degrees, alt.degrees) self.display.set_north(self.pointer.declination) g = self.pointer.compass.gravity self.display.set_gravity(g[0], g[1], g[2]) Avalon.info("ISS Position Update:") #print(self.pointer.azimuth) #print(self.pointer) print( 'Elevation :{0:6.3f}\tAzimuth :{1:6.3f}\nArm Correction:{3:6.3f}\tBase Correction:{2:6.3f}' .format(alt.degrees, az.degrees, float(self.pointer.base_correction), float(self.pointer.arm_correction))) print("Elev: {1:6.3f}\tAz: {0:6.3f}".format( float(self.pointer.azimuth), float(self.pointer.elevation))) print(self.pointer.compass) self.motor_base.set_azimuth(float(self.pointer.azimuth)) self.motor_arm.set_azimuth(float(self.pointer.elevation)) time.sleep(2.5)
def print_peer_config(peer): """ Print the configuration of a specific peer Input takes one Peer object. """ Avalon.info('Peer {} information summary:'.format(peer.address)) if peer.address: print('Address: {}'.format(peer.address)) if peer.public_address: print('Public Address: {}'.format(peer.public_address)) if peer.listen_port: print('Listen Port: {}'.format(peer.listen_port)) print('Private Key: {}'.format(peer.private_key)) if peer.keep_alive: print('Keep Alive: {}'.format(peer.keep_alive))
def check_model_type(args): """ Check if the model demanded from cli argument is legal. """ models_available = [ 'upconv_7_anime_style_art_rgb', 'upconv_7_photo', 'anime_style_art_rgb', 'photo', 'anime_style_art_y' ] if args.model_type not in models_available: Avalon.error('Specified model type not found!') Avalon.info('Available models:') for model in models_available: print(model) exit(1)
def multiThreadGetSubPage(self, threadNumber: int = 16, expectedPageSize: int = 30): workQueue = queue.Queue(threadNumber) threadLock = threading.Lock() exitFlag = False threadList = [] def subPageThread(): while not exitFlag: threadLock.acquire() if not workQueue.empty(): getArgs = workQueue.get() threadLock.release() self.__getSubPageBehavior(*getArgs) else: threadLock.release() time.sleep(1) for i in range(threadNumber): threadName = 'SubFloorThread#%s' % i newThread = threading.Thread(target=subPageThread) newThread.setName(threadName) newThread.start() threadList.append(newThread) totalFloorNumber = self.__db.getlastFloorNum() Avalon.debug_info('[%s]Start Read Replies,Total %s Floor' % (self.__tid, totalFloorNumber)) for floorNum in range(totalFloorNumber): dbResult = self.__db.checkExistFloor(floorNum + 1) if not dbResult: continue replyID = dbResult[1] replyNumber = dbResult[2] replyPageNumber = self.__calcPageNum(expectedPageSize, replyNumber) for i in range(replyPageNumber): workQueue.put((replyID, i + 1)) while not workQueue.empty(): time.sleep(1) exitFlag = 1 for i in threadList: i.join() Avalon.info('[%s] Get Sub Floor Page Success' % self.__tid)
def _get_arp_controller_driver(self): """ Choose ARP controller driver """ print(Avalon.FM.BD + '\nConfigure ARP Controller Driver' + Avalon.FM.RST) # Inform the user which driver is active on current system if shutil.which('nft'): Avalon.info('nftables is available') if shutil.which('arptables'): Avalon.info('arptables is available') while True: driver = Avalon.gets('Please choose an ARP controller driver (nftables/arptables): ') if driver == 'nftables' or driver == 'arptables': self.config['ArpController']['driver'] = driver break else: Avalon.error('Invalid ARP controller driver chosen')
def print_peer_config(peer): """ Print the configuration of a specific peer Input takes one Peer object. """ if peer.alias: Avalon.info(f'{peer.alias} information summary:') else: Avalon.info(f'{peer.address} information summary:') if peer.description: print(f'Description: {peer.description}') if peer.address: print(f'Address: {peer.address}') if peer.public_address: print(f'Public Address: {peer.public_address}') if peer.listen_port: print(f'Listen Port: {peer.listen_port}') print(f'Private Key: {peer.private_key}') if peer.keep_alive: print(f'Keep Alive: {peer.keep_alive}')
def _install_scutum_files(self): """ Install all SCUTUM files into system """ print(Avalon.FM.BD + 'Choose Installation Directory (Enter for default)' + Avalon.FM.RST) installation_dir = Avalon.gets('Choose Installation Path (\"/usr/share/scutum\"):') if installation_dir.strip(' ') != '' and installation_dir[-1] == '/': self.INSTALL_DIR = installation_dir[0:-1] # strip last '/' if exists. breaks program path format Avalon.info(f'Changed installation directory to: {Avalon.FM.BD}{self.INSTALL_DIR}{Avalon.FM.RST}') elif installation_dir.strip(' ') != '': self.INSTALL_DIR = installation_dir Avalon.info(f'Changed installation directory to: {Avalon.FM.BD}{self.INSTALL_DIR}{Avalon.FM.RST}') else: Avalon.info(f'Using default installation directory: {Avalon.FM.BD}{self.INSTALL_DIR}{Avalon.FM.RST}') # If files are not already in place if self.INSTALLER_DIR != self.INSTALL_DIR: if os.path.isdir(self.INSTALL_DIR): shutil.rmtree(self.INSTALL_DIR) # delete existing old scutum files shutil.copytree(self.INSTALLER_DIR, self.INSTALL_DIR) # Remove executable in PATH if exists if os.path.islink(self.SCUTUM_BIN_FILE) or os.path.isfile(self.SCUTUM_BIN_FILE): os.remove(self.SCUTUM_BIN_FILE) # Remove old file or symbolic links # Link scutum main executable to PATH # This will allow user to use the "scutum" command Utilities.execute(['ln', '-s', f'{self.INSTALL_DIR}/bin/scutum.py', self.SCUTUM_BIN_FILE])
def initialize(self, purge=True): """ Checks and adjusts the default rules of ufw which control outgoing data and incoming data. We drop all incoming data by default This will only be ran when scutum is being installed """ if purge: Utilities.execute(['ufw', '--force', 'reset']) # Reset ufw configurations Utilities.execute(['rm', '-f', '/etc/ufw/*.*.*']) # Delete automatic backups coutparsed = Utilities.execute(['ufw', 'status', 'verbose']) for line in coutparsed: if 'Default:' in line: if not (line.split(' ')[1] + line.split(' ')[2] == 'deny(incoming),'): print(line.split(' ')[1] + line.split(' ')[2]) Avalon.info('[UFW]: Adjusting default rule for incoming packages to drop\n') Utilities.execute(['ufw', 'default', 'deny', 'incoming']) if not (line.split(' ')[3] + line.split(' ')[4] == 'allow(outgoing),'): line.split(' ')[3] + line.split(' ')[4] Avalon.info('[UFW]: Adjusting default rule for outgoing packages to allow\n') Utilities.execute(['ufw', 'default', 'allow', 'outgoing']) if 'inactive' in line: Avalon.info('[UFW]: Enabling ufw\n') Utilities.execute(['ufw', 'enable'])
def multiThreadGetMain(self, threadNumber: int = 8): workQueue = queue.Queue(threadNumber) threadLock = threading.Lock() exitFlag = False threadList = [] def mainFloorThread(): while not exitFlag: threadLock.acquire() if not workQueue.empty(): pageNumber = workQueue.get() threadLock.release() self.__getPostBehavior(pageNumber) else: threadLock.release() time.sleep(1) for i in range(threadNumber): threadName = 'PostThread #%s' % i newThread = threading.Thread(target=mainFloorThread) newThread.setName(threadName) newThread.start() threadList.append(newThread) self.__getPostBehavior(1) dbRead = self.__db.checkExistPage(1)[1] if not dbRead: Avalon.critical('Can\'t Get Page 1,Program Exit!') quit(1) totalPages = int(json.loads(dbRead)['page']['total_page']) for i in range(totalPages): workQueue.put(i + 1) while not workQueue.empty(): time.sleep(1) exitFlag = True for i in threadList: i.join() Avalon.info('[%s]Get All Pages Success' % self.__tid)
def upgrade_kpm(): """ upgrade KPM Upgrades KPM by downloading the latest version from GitHub. Replaces the current file being executed. """ try: # get python script web page kpm_request = requests.get(GITHUB_KPM_FILE) if kpm_request.status_code != requests.codes.ok: kpm_request.raise_for_status() # write web page content to file with open(os.path.abspath(__file__), 'wb') as kpm_file: kpm_file.write(kpm_request.content) Avalon.info('KPM has been updated successfully') Avalon.info('Please relaunch KPM') except Exception: Avalon.error('There was an error updating KPM') Avalon.warning('You might have to reinstall KPM manually')
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 start(self): """ Dev: K4YT3X IZAYOI Date Created: Jan 15, 2018 Last Modified: Jan 16, 2018 Dev: Reimannsum Last Modified: Aug 27, 2019 This method is the main ISS pointer controller it runs infinitively until Ctrl^C is pressed. """ ts = load.timescale() stations_url = 'http://celestrak.com/NORAD/elements/stations.txt' satellites = load.tle(stations_url) satellite = satellites['ISS (ZARYA)'] observer = Topos('42.5337N', '83.7384W') while True: t = ts.now() days = t - satellite.epoch if abs(days) > 14: satellites = load.tle(stations_url, reload=True) satellite = satellites['ISS (ZARYA)'] difference = satellite - observer topocentric = difference.at(t) alt, az, distance = topocentric.altaz() # unless the radians would be more use useful elevation = alt.degrees # unless the radians would be more use useful direction = az.degrees Avalon.info("ISS Position Update:") print('Elevation :{}\nAzimuth :{}\n'.format(elevation, direction)) self.motor.set_azimuth(float(direction)) self.servo.set_angle(float(elevation)) time.sleep(5)
def get_proxies() -> collections.deque: """get a list of proxies from ProxyScrape Returns: collections.deque: a deque of retrieved proxies """ proxies_request = requests.get( # "https://api.proxyscrape.com/v2/?request=getproxies&protocol=http&timeout=10000&country=all&ssl=all&anonymity=all" "https://api.proxyscrape.com/v2/?request=getproxies&protocol=socks4&timeout=10000&country=all" ) # if response status code is 200, return the list of retrieved proxies if proxies_request.status_code == requests.codes.ok: proxies = proxies_request.text.split() Avalon.info( f"Successfully retried a list of proxies {Avalon.FM.RST}({len(proxies)} proxies)" ) return collections.deque(proxies) # requests failed to download the list of proxies, raise an exception else: Avalon.error("An error occured while retrieving a list of proxies") proxies_request.raise_for_status()
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')
def upgrade_kpm(): """ upgrade KPM Upgrades KPM by downloading the latest version from GitHub. Replaces the current file being executed. """ try: latest_json = requests.get( "https://api.github.com/repos/k4yt3x/kpm/releases/latest").json() for asset in latest_json["assets"]: if asset["name"] == "kpm.py": latest_version_url = asset["browser_download_url"] break else: Avalon.warning("Unable to find the latest version's download URL") return # get python script web page kpm_request = requests.get(latest_version_url) if kpm_request.status_code != requests.codes.ok: kpm_request.raise_for_status() # write web page content to file with pathlib.Path(__file__).open(mode="wb") as kpm_file: kpm_file.write(kpm_request.content) Avalon.info("KPM has been updated successfully") Avalon.info("Please relaunch KPM") sys.exit(0) except Exception as e: Avalon.error("There was an error updating KPM") Avalon.warning("You might have to reinstall KPM manually") raise e
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 run(self): """ Main controller for Video2X This function controls the flow of video conversion and handles all necessary functions. """ # parse arguments for waifu2x # check argument sanity self._check_arguments() # convert paths to absolute paths self.input_video = self.input_video.absolute() self.output_video = self.output_video.absolute() # initialize objects for ffmpeg and waifu2x-caffe fm = Ffmpeg(self.ffmpeg_settings, self.image_format) # extract frames from video fm.extract_frames(self.input_video, self.extracted_frames) Avalon.info('Reading video information') video_info = fm.get_video_info(self.input_video) # analyze original video with ffprobe and retrieve framerate # width, height = info['streams'][0]['width'], info['streams'][0]['height'] # find index of video stream video_stream_index = None for stream in video_info['streams']: if stream['codec_type'] == 'video': video_stream_index = stream['index'] break # exit if no video stream found if video_stream_index is None: Avalon.error('Aborting: No video stream found') raise StreamNotFoundError('no video stream found') # get average frame rate of video stream framerate = float( Fraction( video_info['streams'][video_stream_index]['avg_frame_rate'])) fm.pixel_format = video_info['streams'][video_stream_index]['pix_fmt'] # get a dict of all pixel formats and corresponding bit depth pixel_formats = fm.get_pixel_formats() # try getting pixel format's corresponding bti depth try: self.bit_depth = pixel_formats[fm.pixel_format] except KeyError: Avalon.error(f'Unsupported pixel format: {fm.pixel_format}') raise UnsupportedPixelError( f'unsupported pixel format {fm.pixel_format}') Avalon.info(f'Framerate: {framerate}') # width/height will be coded width/height x upscale factor if self.scale_ratio: original_width = video_info['streams'][video_stream_index]['width'] original_height = video_info['streams'][video_stream_index][ 'height'] self.scale_width = int(self.scale_ratio * original_width) self.scale_height = int(self.scale_ratio * original_height) # upscale images one by one using waifu2x Avalon.info('Starting to upscale extracted images') self._upscale_frames() Avalon.info('Upscaling completed') # frames to Video Avalon.info('Converting extracted frames into video') # use user defined output size fm.convert_video(framerate, f'{self.scale_width}x{self.scale_height}', self.upscaled_frames) Avalon.info('Conversion completed') # migrate audio tracks and subtitles Avalon.info('Migrating audio tracks and subtitles to upscaled video') fm.migrate_audio_tracks_subtitles(self.input_video, self.output_video, self.upscaled_frames)
except KeyError: self.send_error(404, 'File Not Found: {}'.format(self.path)) def read_redirect_table(path): """ read redirect table content into dictionary Arguments: path {string} -- path to redirect table Returns: dictionary -- dictionary of redirecting rules """ with open(path, 'r') as redirect_table: return json.loads(redirect_table.read()) # parse command line arguments args = process_arguments() # load redirecting rules redirect_table = read_redirect_table(args.redirect_table) # setup socketserver socketserver.TCPServer.allow_reuse_address = True handler = socketserver.TCPServer((args.bind, args.port), request_handler) # start socket server Avalon.info('Starting HTTP server') handler.serve_forever()
# start execution try: # start timer begin_time = time.time() # initialize upscaler object upscaler = Upscaler(input_path=video2x_args.input, output_path=video2x_args.output, driver_settings=driver_settings, ffmpeg_settings=ffmpeg_settings, gifski_settings=gifski_settings) # set upscaler optional options upscaler.driver = video2x_args.driver upscaler.scale_ratio = video2x_args.ratio upscaler.processes = video2x_args.processes upscaler.video2x_cache_directory = video2x_cache_directory upscaler.image_format = image_format upscaler.preserve_frames = preserve_frames # run upscaler upscaler.run() Avalon.info( _('Program completed, taking {} seconds').format( round((time.time() - begin_time), 5))) except Exception: Avalon.error(_('An exception has occurred')) traceback.print_exc()
def run(self): """Main controller for Video2X This function controls the flow of video conversion and handles all necessary functions. """ # external stop signal when called in a thread self.running = True # define process pool to contain processes self.process_pool = [] # load driver modules DriverWrapperMain = getattr( importlib.import_module(f"wrappers.{self.driver}"), "WrapperMain") self.driver_object = DriverWrapperMain(self.driver_settings) # load options from upscaler class into driver settings self.driver_object.load_configurations(self) # initialize FFmpeg object self.ffmpeg_object = Ffmpeg( self.ffmpeg_settings, extracted_frame_format=self.extracted_frame_format) # define processing queue self.processing_queue = queue.Queue() Avalon.info(_("Loading files into processing queue")) Avalon.debug_info(_("Input path(s): {}").format(self.input)) # make output directory if the input is a list or a directory if isinstance(self.input, list) or self.input.is_dir(): self.output.mkdir(parents=True, exist_ok=True) input_files = [] # if input is single directory # put it in a list for compability with the following code if not isinstance(self.input, list): input_paths = [self.input] else: input_paths = self.input # flatten directories into file paths for input_path in input_paths: # if the input path is a single file # add the file's path object to input_files if input_path.is_file(): input_files.append(input_path) # if the input path is a directory # add all files under the directory into the input_files (non-recursive) elif input_path.is_dir(): input_files.extend( [f for f in input_path.iterdir() if f.is_file()]) output_paths = [] for input_path in input_files: # get file type # try python-magic if it's available try: input_file_mime_type = magic.from_file(str( input_path.absolute()), mime=True) input_file_type = input_file_mime_type.split("/")[0] input_file_subtype = input_file_mime_type.split("/")[1] except Exception: input_file_mime_type = input_file_type = input_file_subtype = "" # if python-magic doesn't determine the file to be an image/video file # fall back to mimetypes to guess the file type based on the extension if input_file_type not in ["image", "video"]: # in case python-magic fails to detect file type # try guessing file mime type with mimetypes input_file_mime_type = mimetypes.guess_type(input_path.name)[0] input_file_type = input_file_mime_type.split("/")[0] input_file_subtype = input_file_mime_type.split("/")[1] Avalon.debug_info( _("File MIME type: {}").format(input_file_mime_type)) # set default output file suffixes # if image type is GIF, default output suffix is also .gif if input_file_mime_type == "image/gif": output_path = self.output / self.output_file_name_format_string.format( original_file_name=input_path.stem, extension=".gif") elif input_file_type == "image": output_path = self.output / self.output_file_name_format_string.format( original_file_name=input_path.stem, extension=self.image_output_extension, ) elif input_file_type == "video": output_path = self.output / self.output_file_name_format_string.format( original_file_name=input_path.stem, extension=self.video_output_extension, ) # if file is none of: image, image/gif, video # skip to the next task else: Avalon.error( _("File {} ({}) neither an image nor a video").format( input_path, input_file_mime_type)) Avalon.warning(_("Skipping this file")) continue # if there is only one input file # do not modify output file suffix if isinstance(self.input, pathlib.Path) and self.input.is_file(): output_path = self.output output_path_id = 0 while str(output_path) in output_paths: output_path = output_path.parent / pathlib.Path( f"{output_path.stem}_{output_path_id}{output_path.suffix}") output_path_id += 1 # record output path output_paths.append(str(output_path)) # push file information into processing queue self.processing_queue.put(( input_path.absolute(), output_path.absolute(), input_file_mime_type, input_file_type, input_file_subtype, )) # check argument sanity before running self._check_arguments() # record file count for external calls self.total_files = self.processing_queue.qsize() Avalon.info(_("Loaded files into processing queue")) # print all files in queue for debugging for job in self.processing_queue.queue: Avalon.debug_info(_("Input file: {}").format(job[0].absolute())) try: while not self.processing_queue.empty(): # get new job from queue ( self.current_input_file, output_path, input_file_mime_type, input_file_type, input_file_subtype, ) = self.processing_queue.get() # get current job starting time for GUI calculations self.current_processing_starting_time = time.time() # get video information JSON using FFprobe Avalon.info(_("Reading file information")) file_info = self.ffmpeg_object.probe_file_info( self.current_input_file) # create temporary directories for storing frames self.create_temp_directories() # start handling input # if input file is a static image if input_file_type == "image" and input_file_subtype != "gif": Avalon.info(_("Starting upscaling image")) # copy original file into the pre-processing directory shutil.copy( self.current_input_file, self.extracted_frames / self.current_input_file.name, ) width = int(file_info["streams"][0]["width"]) height = int(file_info["streams"][0]["height"]) framerate = self.total_frames = 1 # elif input_file_mime_type == 'image/gif' or input_file_type == 'video': else: Avalon.info(_("Starting upscaling video/GIF")) # find index of video stream video_stream_index = None for stream in file_info["streams"]: if stream["codec_type"] == "video": video_stream_index = stream["index"] break # exit if no video stream found if video_stream_index is None: Avalon.error(_("Aborting: No video stream found")) raise StreamNotFoundError("no video stream found") # get average frame rate of video stream framerate = float( Fraction(file_info["streams"][video_stream_index] ["r_frame_rate"])) width = int( file_info["streams"][video_stream_index]["width"]) height = int( file_info["streams"][video_stream_index]["height"]) # get total number of frames Avalon.info( _("Getting total number of frames in the file")) # if container stores total number of frames in nb_frames, fetch it directly if "nb_frames" in file_info["streams"][video_stream_index]: self.total_frames = int( file_info["streams"][video_stream_index] ["nb_frames"]) # otherwise call FFprobe to count the total number of frames else: self.total_frames = self.ffmpeg_object.get_number_of_frames( self.current_input_file, video_stream_index) # calculate scale width/height/ratio and scaling jobs if required Avalon.info(_("Calculating scaling parameters")) # create a local copy of the global output settings output_scale = self.scale_ratio output_width = self.scale_width output_height = self.scale_height # calculate output width and height if scale ratio is specified if output_scale is not None: output_width = int( math.ceil(width * output_scale / 2.0) * 2) output_height = int( math.ceil(height * output_scale / 2.0) * 2) else: # scale keeping aspect ratio is only one of width/height is given if output_width == 0 or output_width is None: output_width = output_height / height * width elif output_height == 0 or output_height is None: output_height = output_width / width * height output_width = int(math.ceil(output_width / 2.0) * 2) output_height = int(math.ceil(output_height / 2.0) * 2) # calculate required minimum scale ratio output_scale = max(output_width / width, output_height / height) # if driver is one of the drivers that doesn't support arbitrary scaling ratio # TODO: more documentations on this block if self.driver in DRIVER_FIXED_SCALING_RATIOS: # select the optimal driver scaling ratio to use supported_scaling_ratios = sorted( DRIVER_FIXED_SCALING_RATIOS[self.driver]) remaining_scaling_ratio = math.ceil(output_scale) self.scaling_jobs = [] # if the scaling ratio is 1.0 # apply the smallest scaling ratio available if remaining_scaling_ratio == 1: self.scaling_jobs.append(supported_scaling_ratios[0]) else: while remaining_scaling_ratio > 1: for ratio in supported_scaling_ratios: if ratio >= remaining_scaling_ratio: self.scaling_jobs.append(ratio) remaining_scaling_ratio /= ratio break else: found = False for i in supported_scaling_ratios: for j in supported_scaling_ratios: if i * j >= remaining_scaling_ratio: self.scaling_jobs.extend([i, j]) remaining_scaling_ratio /= i * j found = True break if found is True: break if found is False: self.scaling_jobs.append( supported_scaling_ratios[-1]) remaining_scaling_ratio /= supported_scaling_ratios[ -1] else: self.scaling_jobs = [output_scale] # print file information Avalon.debug_info(_("Framerate: {}").format(framerate)) Avalon.debug_info(_("Width: {}").format(width)) Avalon.debug_info(_("Height: {}").format(height)) Avalon.debug_info( _("Total number of frames: {}").format(self.total_frames)) Avalon.debug_info(_("Output width: {}").format(output_width)) Avalon.debug_info(_("Output height: {}").format(output_height)) Avalon.debug_info( _("Required scale ratio: {}").format(output_scale)) Avalon.debug_info( _("Upscaling jobs queue: {}").format(self.scaling_jobs)) # extract frames from video if input_file_mime_type == "image/gif" or input_file_type == "video": self.process_pool.append( (self.ffmpeg_object.extract_frames( self.current_input_file, self.extracted_frames))) self._wait() # if driver is waifu2x-caffe # pass pixel format output depth information if self.driver == "waifu2x_caffe": # get a dict of all pixel formats and corresponding bit depth pixel_formats = self.ffmpeg_object.get_pixel_formats() # try getting pixel format's corresponding bti depth try: self.driver_settings["output_depth"] = pixel_formats[ self.ffmpeg_object.pixel_format] except KeyError: Avalon.error( _("Unsupported pixel format: {}").format( self.ffmpeg_object.pixel_format)) raise UnsupportedPixelError( f"unsupported pixel format {self.ffmpeg_object.pixel_format}" ) # upscale images one by one using waifu2x Avalon.info(_("Starting to upscale extracted frames")) upscale_begin_time = time.time() self.current_pass = 1 if self.driver == "waifu2x_caffe": self.driver_object.set_scale_resolution( output_width, output_height) else: self.driver_object.set_scale_ratio(self.scaling_jobs[0]) self._upscale_frames(self.extracted_frames, self.upscaled_frames) for job in self.scaling_jobs[1:]: self.current_pass += 1 self.driver_object.set_scale_ratio(job) shutil.rmtree(self.extracted_frames) shutil.move(self.upscaled_frames, self.extracted_frames) self.upscaled_frames.mkdir(parents=True, exist_ok=True) self._upscale_frames(self.extracted_frames, self.upscaled_frames) Avalon.info(_("Upscaling completed")) Avalon.info( _("Average processing speed: {} seconds per frame").format( self.total_frames / (time.time() - upscale_begin_time))) # downscale frames with Lanczos Avalon.info(_("Lanczos downscaling frames")) shutil.rmtree(self.extracted_frames) shutil.move(self.upscaled_frames, self.extracted_frames) self.upscaled_frames.mkdir(parents=True, exist_ok=True) for image in tqdm( [ i for i in self.extracted_frames.iterdir() if i.is_file() and i.name.endswith(self.extracted_frame_format) ], ascii=True, desc=_("Downscaling"), ): image_object = Image.open(image) # if the image dimensions are not equal to the output size # resize the image using Lanczos if (image_object.width, image_object.height) != ( output_width, output_height, ): image_object.resize( (output_width, output_height), Image.LANCZOS).save( self.upscaled_frames / image.name) image_object.close() # if the image's dimensions are already equal to the output size # move image to the finished directory else: image_object.close() shutil.move(image, self.upscaled_frames / image.name) # start handling output # output can be either GIF or video if input_file_type == "image" and input_file_subtype != "gif": Avalon.info(_("Exporting image")) # there should be only one image in the directory shutil.move( [ f for f in self.upscaled_frames.iterdir() if f.is_file() ][0], output_path, ) # elif input_file_mime_type == 'image/gif' or input_file_type == 'video': else: # if the desired output is gif file if output_path.suffix.lower() == ".gif": Avalon.info( _("Converting extracted frames into GIF image")) gifski_object = Gifski(self.gifski_settings) self.process_pool.append( gifski_object.make_gif( self.upscaled_frames, output_path, framerate, self.extracted_frame_format, output_width, output_height, )) self._wait() Avalon.info(_("Conversion completed")) # if the desired output is video else: # frames to video Avalon.info( _("Converting extracted frames into video")) self.process_pool.append( self.ffmpeg_object.assemble_video( framerate, self.upscaled_frames)) # f'{scale_width}x{scale_height}' self._wait() Avalon.info(_("Conversion completed")) try: # migrate audio tracks and subtitles Avalon.info( _("Migrating audio, subtitles and other streams to upscaled video" )) self.process_pool.append( self.ffmpeg_object.migrate_streams( self.current_input_file, output_path, self.upscaled_frames, )) self._wait() # if failed to copy streams # use file with only video stream except subprocess.CalledProcessError: traceback.print_exc() Avalon.error(_("Failed to migrate streams")) Avalon.warning( _("Trying to output video without additional streams" )) if input_file_mime_type == "image/gif": # copy will overwrite destination content if exists shutil.copy( self.upscaled_frames / self.ffmpeg_object.intermediate_file_name, output_path, ) else: # construct output file path output_file_name = f"{output_path.stem}{self.ffmpeg_object.intermediate_file_name.suffix}" output_video_path = (output_path.parent / output_file_name) # if output file already exists # create temporary directory in output folder # temporary directories generated by tempfile are guaranteed to be unique # and won't conflict with other files if output_video_path.exists(): Avalon.error(_("Output video file exists")) temporary_directory = pathlib.Path( tempfile.mkdtemp( dir=output_path.parent)) output_video_path = (temporary_directory / output_file_name) Avalon.info( _("Created temporary directory to contain file" )) # move file to new destination Avalon.info( _("Writing intermediate file to: {}"). format(output_video_path.absolute())) shutil.move( self.upscaled_frames / self.ffmpeg_object.intermediate_file_name, output_video_path, ) # increment total number of files processed self.cleanup_temp_directories() self.processing_queue.task_done() self.total_processed += 1 except (Exception, KeyboardInterrupt, SystemExit) as e: with contextlib.suppress(ValueError, AttributeError): self.cleanup_temp_directories() self.running = False raise e # signal upscaling completion self.running = False