def methodPost(url: str, datasEncoded: bytes, headers: dict = {}, maxTryTimes: int = 10): encodedUrl = parse.urlparse(url=url).geturl() requsetMaker = request.Request(url=encodedUrl, headers=headers, data=datasEncoded) for tryTimes in range(maxTryTimes): try: dataPost = request.urlopen(url=requsetMaker, timeout=10).read() except error.URLError as e: Avalon.debug( 'POST "%s" Failed,reason:%s,Program will Try %d Times Later.' % (url, e.reason, maxTryTimes - tryTimes)) except timeout as e: Avalon.debug( 'POST "%s" Timeout,reason:%s,Program will Try %d Times Later.' % (url, e, maxTryTimes - tryTimes)) except: Avalon.debug( 'POST "%s" Timeout with unknown reason,Program will Try %d Times Later.' % (url, maxTryTimes - tryTimes)) else: break else: Avalon.error( 'POST "%s" Failed during all request,please check your network status.' % url) __quitCheck(1) return bytes(dataPost)
def convDataToPerFloor(self): dbGot = json.loads(self.__db.checkExistPage(1)[1]) totalPage = int(dbGot['page']['total_page']) for pageNum in range(totalPage): gotData = self.__db.checkExistPage(pageNum + 1) if not gotData: Avalon.error('Can\'t Get Page %s,Skip' % pageNum) continue gotData = json.loads(gotData[1]) self.__writePageUserNameToDatabase(gotData) for i in gotData['post_list']: replyID = int(i['id']) replyNum = int(i.get('sub_post_number', 0)) floorNumber = int(i['floor']) publishTime = int(i['time']) userID = int(i['author_id']) userName = self.__getUserName(userID) context = str(json.dumps(i)) self.__db.writeFloor(floorNumber, replyID, replyNum, publishTime, userName, context) Avalon.debug_info( '[%s]Floor Info at Page %s Finished.Database Changed %s Record' % (self.__tid, pageNum + 1, self.__dbTotalChange))
def create_temp_directories(self): """create temporary directories """ # if cache directory unspecified, use %TEMP%\video2x if self.video2x_cache_directory is None: self.video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x' # if specified cache path exists and isn't a directory if self.video2x_cache_directory.exists() and not self.video2x_cache_directory.is_dir(): Avalon.error(_('Specified or default cache directory is a file/link')) raise FileExistsError('Specified or default cache directory is a file/link') # if cache directory doesn't exist, try creating it if not self.video2x_cache_directory.exists(): try: Avalon.debug_info(_('Creating cache directory {}').format(self.video2x_cache_directory)) self.video2x_cache_directory.mkdir(parents=True, exist_ok=True) except Exception as exception: Avalon.error(_('Unable to create {}').format(self.video2x_cache_directory)) raise exception # create temp directories for extracted frames and upscaled frames self.extracted_frames = pathlib.Path(tempfile.mkdtemp(dir=self.video2x_cache_directory)) Avalon.debug_info(_('Extracted frames are being saved to: {}').format(self.extracted_frames)) self.upscaled_frames = pathlib.Path(tempfile.mkdtemp(dir=self.video2x_cache_directory)) Avalon.debug_info(_('Upscaled frames are being saved to: {}').format(self.upscaled_frames))
def read_config(): """ Parses configuration This function parses the configuration file and load the configurations into the program TODO: Do something about KeyError """ if not os.path.isfile(args.config): # Configuration not found Avalon.error('SCUTUM configuration file not found! Please re-install SCUTUM!') Avalon.warning('Please run \"scutum --install\" before using it for the first time') raise FileNotFoundError(args.config) # Initialize python confparser and read config with open(args.config, 'r') as raw_config: config = json.load(raw_config) # Get controlled interfaces interfaces = [] for interface in config['Interfaces']['interfaces']: if os.path.isdir(f'/sys/class/net/{interface}'): # Check if interface is connected interfaces.append(interface) # Get controlled network controllers network_controllers = config['NetworkControllers']['controllers'] # Check if we should handle ufw ufw_handled = config['Ufw']['handled'] # Get ARP Controller driver arp_driver = config['ArpController']['driver'] return interfaces, network_controllers, ufw_handled, arp_driver
def install_packages(packages): """ Install a package using system package manager This method is currently using os.system instead of subprocess.run or subprocess.Popen because subprocess doesn't seem to handle some of the TUIs well. """ Avalon.warning( 'If the installation is unsuccessful, you should consider updating the package manager cache', log=False) # Convert packages list into a string if len(packages) > 1: packages_string = ' '.join(packages) else: packages_string = packages[0] if shutil.which('apt-get'): return os.system(f'apt-get install {packages_string} -y') elif shutil.which('yum'): return os.system(f'yum install {packages_string} -y') elif shutil.which('pacman'): return os.system(f'pacman -S {packages_string} --noconfirm') else: Avalon.error( 'Sorry, we can\'t find a package manager that we currently support. Aborting..' ) Avalon.error('Currently Supported: apt, yum, pacman') return False
def get_path(text): """ Get path and validate """ while True: path = Avalon.gets(text) if os.path.isdir(path): return path Avalon.error('{} id not a directory/folder'.format(path))
def getPost(self: None, pageNumber: int, ajax: bool = True, useTemp: bool = True): # 获得html源文件函数 self.__workPageNumber = pageNumber link = self.__postLink + str(pageNumber) existTemp = self.__tempSave.getSameTemp() if existTemp.get('html') and useTemp: for i in existTemp['html']: if int(i[1]) == int(pageNumber): Avalon.debug_info('第%d页已经在临时文件中存在,跳过' % pageNumber) return self.__tempSave.readFileByID(i) if ajax is False: link = link.replace('ajax=1&', '') for tryTimes in range(1, 11): try: postRequest = request.Request(link) try: # 设置程序请求头,伪装爬虫(必要性存疑) postRequest.add_header('User-Agent', (random.choice( self.__userAgent)).replace('\n', '')) postRequest.add_header('Referer', 'https://tieba.baidu.com') except: continue else: postRead: bytes = request.urlopen(postRequest, timeout=5).read() if self.debug: Avalon.debug_info('链接:"%s"请求头:%s.' % (link, postRequest.headers)) # 错误处理 except error.URLError as e: Avalon.warning("获取帖子正文失败!原因:%s(%s/10)" % (str(e.reason), str(tryTimes))) except timeout as e: Avalon.warning("获取帖子正文失败!原因:%s(%s/10)" % (str(e), str(tryTimes))) except KeyboardInterrupt: Avalon.critical("用户强制退出") quit(1) except: Avalon.warning("获取帖子正文失败!原因:未知错误(%s/10)" % tryTimes) # 没有错误,结束循环 else: if self.debug: Avalon.debug_info('Link %s Get Successed.' % link) break else: Avalon.error('获取失败!') if self.debug: Avalon.debug('Link:%s' % link) quit(1) if useTemp is True: self.__tempSave.savePostRaw(postRead.decode(errors='ignore'), pageNumber=pageNumber) return (postRead.decode(errors='ignore'))
def command_interpreter(commands): """ WGC shell command interpreter This function interprets commands from CLI or the interactive shell, and passes the parameters to the corresponding functions. """ 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() == 'showpeers': for peer in pm.peers: print_peer_config(peer) result = 0 elif commands[1].lower() == 'jsonloadprofile': result = pm.json_load_profile(commands[2]) elif commands[1].lower() == 'jsonsaveprofile': result = pm.json_save_profile(commands[2]) elif commands[1].lower() == 'pickleloadprofile': result = pm.pickle_load_profile(commands[2]) elif commands[1].lower() == 'picklesaveprofile': result = pm.pickle_save_profile(commands[2]) elif commands[1].lower() == 'newprofile': result = pm.new_profile() elif commands[1].lower() == 'addpeer': result = add_peer() elif commands[1].lower() == 'deletepeer': result = delete_peer(commands[2]) elif commands[1].lower() == 'generateconfigs': result = generate_configs(commands[2]) elif commands[1].lower() == 'exit' or commands[1].lower() == 'quit': Avalon.warning('Exiting') exit(0) elif len(possibilities) > 0: Avalon.warning(f'Ambiguous command \"{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 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 __getReplyIDList(self, pageNumber: int): dbResult = self.__db.checkExistPage(pageNumber) if not dbResult: Avalon.error('Failed To Get Page %s In Database' % pageNumber) return False dbResultDecode = json.loads(dbResult) replyIDList = [] for perFloor in dbResultDecode['post_list']: replyID = perFloor.get('id') if replyID: replyIDList.append(int(replyID)) return replyIDList
def _wait(self): """wait for subprocesses in process pool to complete""" Avalon.debug_info(_("Main process waiting for subprocesses to exit")) try: # while process pool not empty while self.process_pool: # if stop signal received, terminate all processes if self.running is False: raise SystemExit for process in self.process_pool: process_status = process.poll() # if process finished if process_status is None: continue # if return code is not 0 elif process_status != 0: Avalon.error( _("Subprocess {} exited with code {}").format( process.pid, process_status ) ) raise subprocess.CalledProcessError( process_status, process.args ) else: Avalon.debug_info( _("Subprocess {} exited with code {}").format( process.pid, process_status ) ) self.process_pool.remove(process) time.sleep(0.1) except (KeyboardInterrupt, SystemExit) as e: Avalon.warning(_("Stop signal received")) self._terminate_subprocesses() raise e except (Exception, subprocess.CalledProcessError) as e: Avalon.error(_("Subprocess execution ran into an error")) self._terminate_subprocesses() raise e
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 __init__(self): """ Keyword Arguments: log {object} -- object of logger (default: {False}) Raises: FileNotFoundError -- raised when UFW not installed """ if shutil.which('ufw') is None: # Detect if ufw installed print(Avalon.FM.BD + Avalon.FG.R + '\nWe have detected that you don\'t have UFW installed!' + Avalon.FM.RST) print('UFW Firewall function requires UFW to run') if not Utilities.install_package('ufw'): Avalon.error('ufw is required for this function. Exiting...') raise FileNotFoundError('File: \"/usr/sbin/ufw\" not found')
def get_psk_length(): """Asks the user for desired password length Returns: int -- length of PSK """ while 1: try: psk_length = int(Avalon.gets("Desired PSK Length: ")) if psk_length < 8 or psk_length > 63: Avalon.error("PSK length should be between 8 and 63") continue break except ValueError: Avalon.error("Invalid Input") return psk_length
def command_interpreter(commands): """ WGC shell command interpreter This function interprets commands from CLI or the interactive shell, and passes the parameters to the corresponding functions. """ 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() == 'gencacert': cacert = CACert(commands[2]) cacert.generate() result = 0 elif commands[1].lower() == 'genusercert': usercert = UserCert(commands[2], commands[3], commands[4]) usercert.generate() result = 0 elif commands[1].lower() == 'genservercert': servercert = ServerCert(commands[2], commands[3], commands[4], commands[5]) servercert.generate() result = 0 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 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 _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 _get_controlled_interfaces(self): """ Get interfaces to be controlled by SCUTUM """ # Get controlled interfaces ifaces_selected = [] ifaces = [] # List all available interfaces with open('/proc/net/dev', 'r') as dev: for line in dev: try: if line.split(':')[1]: ifaces.append(line.split(':')[0]) except IndexError: pass # Enroll controlled interfaces while True: print(Avalon.FM.BD + '\nWhich interface do you want scutum to control?' + Avalon.FM.RST) if not len(ifaces) == 0: index = 0 for iface in ifaces: if iface.replace(' ', '') not in ifaces_selected: print(f'{str(index)}. {iface.replace(" ", "")}') index += 1 print('x. Manually Enter') print(Avalon.FM.BD + 'Press [ENTER] when complete' + Avalon.FM.RST) selection = Avalon.gets('Please select (index number): ') try: if selection == 'x': manif = Avalon.gets('Interface: ') if manif not in ifaces_selected: ifaces_selected.append(manif) elif selection == '': if len(ifaces_selected) != 0: break else: Avalon.error('You have not selected any interfaces yet') elif int(selection) >= len(ifaces): Avalon.error('Selected interface doesn\'t exist!') else: ifaces_selected.append(ifaces[int(selection)].replace(' ', '')) except ValueError: Avalon.error('Invalid Input!') Avalon.error('Please enter the index number!') # Put controlled interfaces into configuraion self.config['Interfaces']['interfaces'] = ifaces_selected
def main(): """ WireGuard Mesh Configurator main function This function controls the main flow of this program. """ try: if sys.argv[1].lower() == 'help': print_help() exit(0) except IndexError: pass # Begin command interpreting try: if sys.argv[1].lower() == 'interactive' or sys.argv[1].lower( ) == 'int': print_welcome() # Set command completer completer = ShellCompleter(COMMANDS) readline.set_completer(completer.complete) readline.parse_and_bind('tab: complete') # Launch interactive trojan shell prompt = '{}[WGC]> {}'.format(Avalon.FM.BD, Avalon.FM.RST) while True: command_interpreter([''] + input(prompt).split(' ')) else: # Return to shell with command return value exit(command_interpreter(sys.argv[0:])) except IndexError: Avalon.warning('No commands specified') print_help() exit(0) except (KeyboardInterrupt, EOFError): Avalon.warning('Exiting') exit(0) except Exception: Avalon.error('Exception caught') traceback.print_exc() exit(1)
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 getPostInfo(self: None): postRaw = self.getPost(pageNumber=1, ajax=False, useTemp=False) postGet = etree.HTML(postRaw) if postGet.xpath('//body[@class="page404"]'): Avalon.error('此贴已经被删除!') raise FileNotFoundError postTitle = postGet.xpath('//div[@class="wrap2"]//h3/@title') postAuthor = postGet.xpath( '//div[@class="p_postlist"]/div[@class][1]//div/@author') postPageNum = postGet.xpath( '//div/div[@id]//div[@id]//li/span[@class="red"][last()]/text()') if not (postTitle and postAuthor and postPageNum): Avalon.critical('程序无法正确获得帖子信息') quit(1) finalInfo = { 'Author': str(postAuthor[0]), 'Title': str(postTitle[0]), 'TotalPage': int(postPageNum[0]) } self.postInfo = finalInfo if self.debug: Avalon.debug_info(self.postInfo) return finalInfo
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 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 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 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 _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 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)
Arguments: config_file {pathlib.Path} -- video2x configuration file pathlib.Path Returns: dict -- dictionary of video2x configuration """ with open(config_file, 'r') as config: return yaml.load(config, Loader=yaml.FullLoader) # /////////////////// Execution /////////////////// # # this is not a library if __name__ != '__main__': Avalon.error(_('This file cannot be imported')) raise ImportError(f'{__file__} cannot be imported') # print video2x logo print_logo() # parse command line arguments video2x_args, driver_args = parse_arguments() # display version and lawful informaition if video2x_args.version: print(LEGAL_INFO) sys.exit(0) # redirect output to both terminal and log file if video2x_args.nolog is False:
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
def _check_arguments(self): if isinstance(self.input, list): if self.output.exists() and not self.output.is_dir(): Avalon.error(_("Input and output path type mismatch")) Avalon.error( _("Input is multiple files but output is not directory")) raise ArgumentError("input output path type mismatch") for input_path in self.input: if not input_path.is_file() and not input_path.is_dir(): Avalon.error( _("Input path {} is neither a file nor a directory"). format(input_path)) raise FileNotFoundError( f"{input_path} is neither file nor directory") with contextlib.suppress(FileNotFoundError): if input_path.samefile(self.output): Avalon.error( _("Input directory and output directory cannot be the same" )) raise FileExistsError( "input directory and output directory are the same" ) # if input is a file elif self.input.is_file(): if self.output.is_dir(): Avalon.error(_("Input and output path type mismatch")) Avalon.error(_("Input is single file but output is directory")) raise ArgumentError("input output path type mismatch") if self.output.suffix == "": Avalon.error(_("No suffix found in output file path")) Avalon.error(_("Suffix must be specified")) raise ArgumentError("no output file suffix specified") # if input is a directory elif self.input.is_dir(): if self.output.is_file(): Avalon.error(_("Input and output path type mismatch")) Avalon.error( _("Input is directory but output is existing single file")) raise ArgumentError("input output path type mismatch") with contextlib.suppress(FileNotFoundError): if self.input.samefile(self.output): Avalon.error( _("Input directory and output directory cannot be the same" )) raise FileExistsError( "input directory and output directory are the same") # if input is neither else: Avalon.error(_("Input path is neither a file nor a directory")) raise FileNotFoundError( f"{self.input} is neither file nor directory") # check FFmpeg settings ffmpeg_path = pathlib.Path(self.ffmpeg_settings["ffmpeg_path"]) if not ((pathlib.Path(ffmpeg_path / "ffmpeg.exe").is_file() and pathlib.Path(ffmpeg_path / "ffprobe.exe").is_file()) or (pathlib.Path(ffmpeg_path / "ffmpeg").is_file() and pathlib.Path(ffmpeg_path / "ffprobe").is_file())): Avalon.error( _("FFmpeg or FFprobe cannot be found under the specified path") ) Avalon.error(_("Please check the configuration file settings")) raise FileNotFoundError(self.ffmpeg_settings["ffmpeg_path"]) # check if driver settings driver_settings = copy.deepcopy(self.driver_settings) driver_path = driver_settings.pop("path") # check if driver path exists if not (pathlib.Path(driver_path).is_file() or pathlib.Path(f"{driver_path}.exe").is_file()): Avalon.error( _("Specified driver executable directory doesn't exist")) Avalon.error(_("Please check the configuration file settings")) raise FileNotFoundError(driver_path) # parse driver arguments using driver's parser # the parser will throw AttributeError if argument doesn't satisfy constraints try: driver_arguments = [] for key in driver_settings.keys(): value = driver_settings[key] if value is None or value is False: continue else: if len(key) == 1: driver_arguments.append(f"-{key}") else: driver_arguments.append(f"--{key}") # true means key is an option if value is not True: driver_arguments.append(str(value)) DriverWrapperMain = getattr( importlib.import_module(f"wrappers.{self.driver}"), "WrapperMain") DriverWrapperMain.parse_arguments(driver_arguments) except AttributeError as e: Avalon.error( _("Failed to parse driver argument: {}").format(e.args[0])) raise e