def restart(self):
avalon.info('Restarting trojan server')
os.system('systemctl restart trojan')
def start(self):
avalon.info('Starting trojan server')
os.system('systemctl start trojan')
def stop(self):
avalon.info('Stopping trojan server')
os.system('systemctl stop trojan')
def install(self):
"""
This is the main function for installer
"""
global ifacesSelected
config = configparser.ConfigParser()
config["Interfaces"] = {}
config["NetworkControllers"] = {}
config["Ufw"] = {}
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
os.system("ln -s " + self.INSTALL_DIR + "/scutum.py " +
self.SCUTUM_BIN_FILE)
self.install_service() # install and register service files
os.system("systemctl enable scutum") # enable service
os.system("systemctl start scutum") # start service
if not os.path.isfile('/usr/bin/arptables') and not os.path.isfile(
'/sbin/arptables'): # Detect if arptables installed
print(
avalon.FM.BD + avalon.FG.R +
'\nWe have detected that you don\'t have arptables installed!'
+ avalon.FM.RST)
print('SCUTUM requires arptables to run')
if not self.sysInstallPackage("arptables"):
avalon.error("arptables is required for scutum. Exiting...")
exit(1)
ifacesSelected = []
while True:
print(avalon.FM.BD +
'\nWhich interface do you wish to install for?' +
avalon.FM.RST)
ifaces = []
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
if not len(ifaces) == 0:
idx = 0
for iface in ifaces:
print(str(idx) + '. ' + iface.replace(' ', ''))
idx += 1
print('99. Manually Enter')
selection = avalon.gets('Please select (index number): ')
try:
if selection == '99':
manif = avalon.gets('Interface: ')
if manif not in ifacesSelected:
ifacesSelected.append(manif)
if avalon.ask('Add more interfaces?', False):
pass
else:
break
elif int(selection) >= len(ifaces):
avalon.error('Selected interface doesn\'t exist!')
else:
ifacesSelected.append(ifaces[int(selection)].replace(
' ', ''))
if avalon.ask('Add more interfaces?', False):
pass
else:
break
except ValueError:
avalon.error('Invalid Input!')
avalon.error('Please enter the index number!')
config["Interfaces"]["interfaces"] = ",".join(ifacesSelected)
while True:
print(avalon.FM.BD +
'\nWhich network controller do you want to install for?' +
avalon.FM.RST)
print('1. WICD')
print('2. Network-Manager')
print('3. Both')
selection = avalon.gets('Please select: (index number): ')
if selection == '1':
if self.installWicdScripts() is not True:
avalon.error(
"SCUTUM Script for WICD has failed to install!")
avalon.error("Aborting Installation...")
exit(1)
config["NetworkControllers"]["controllers"] = "wicd"
break
elif selection == '2':
if self.installNMScripts(ifacesSelected) is not True:
avalon.error(
"SCUTUM Script for NetworkManager has failed to install!"
)
avalon.error("Aborting Installation...")
exit(1)
config["NetworkControllers"]["controllers"] = "NetworkManager"
break
elif selection == '3':
ifaces = ["wicd", "NetworkManager"]
if self.installWicdScripts() is not True:
avalon.warning("Deselected WICD from installation")
ifaces.remove("wicd")
if self.installNMScripts(ifacesSelected) is not True:
avalon.warning(
"Deselected NetworkManager from installation")
ifaces.remove("NetworkManager")
if len(ifaces) == 0:
avalon.error("All SCUTUM Scripts have failed to install!")
avalon.error("Aborting Installation...")
exit(1)
config["NetworkControllers"]["controllers"] = ",".join(ifaces)
break
else:
avalon.error('Invalid Input!')
print(avalon.FM.BD + '\nEnable UFW firewall?' + avalon.FM.RST)
print(
"Do you want SCUTUM to help configuring and enabling UFW firewall?"
)
print("This will prevent a lot of scanning and attacks")
if avalon.ask('Enable?', True):
ufwctrl = Ufw(False)
print("UFW 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):
config["Ufw"]["handled"] = "true"
else:
config["Ufw"]["handled"] = "false"
else:
avalon.info("You can turn it on whenever you change your mind")
print(avalon.FM.BD + '\nInstall Easy TCP controllers?' + avalon.FM.RST)
print("Easy tcp controller helps you open/close ports quickly")
print("ex. \"openport 80\" opens port 80")
print("ex. \"closeport 80\" closes port 80")
print("ex. \"openport 80 443\" opens port 80 and 443")
print("ex. \"closeport 80 443\" closes port 80 and 443")
if avalon.ask("Install Easy TCP conrollers?", True):
self.install_easytcp_controllers()
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):
self.install_scutum_gui()
with open(self.CONFPATH, 'w') as configfile:
config.write(configfile) # Writes configurations
if args.version: # prints program legal / dev / version info
print("Current Version: " + VERSION)
print("Author: K4YT3X")
print("License: GNU GPL v3")
print("Github Page: https://github.com/K4YT3X/scutum")
print("Contact: [email protected]\n")
exit(0)
elif args.status:
"""
Asks systemd-sysv-install if scutum is enabled
by systemctl. May not apply to non-Debian distros
"""
if os.system("/lib/systemd/systemd-sysv-install is-enabled scutum"):
avalon.info("{}SCUTUM is {}{}{}\n".format(avalon.FM.RST, avalon.FG.R,
"NOT ENABLED",
avalon.FM.RST))
else:
avalon.info("{}SCUTUM is {}{}{}\n".format(avalon.FM.RST, avalon.FG.G,
"ENABLED", avalon.FM.RST))
exit(0)
elif os.getuid() != 0: # Multiple components require root access
avalon.error('SCUTUM must be run as root!')
print(avalon.FG.LGR + 'It requires root privilege to operate the system' +
avalon.FM.RST)
exit(1)
exit_code = 0
log = open(LOGPATH, 'a+')
installer = Installer(CONFPATH, log=log)
print('Author: K4YT3X')
print('License: GNU GPL v3')
print('Github Page: https://github.com/K4YT3X/KonaDL')
print('Contact: [email protected]\n')
exit(0)
kona.storage = check_storage_dir(args)
if kona.storage is False:
avalon.error('Please specify storage directory\n')
exit(1)
# If progress file exists
# Ask user if he or she wants to load it
load_progress = False
if kona.progress_files_present():
avalon.info('Progress file found')
if avalon.ask('Continue from where you left off?', True):
kona.load_progress = True
load_progress = True
else:
avalon.info('Creating new download profile')
# Pass terminal arguments to libkonadl object
kona.yandere = args.yandere
kona.safe = args.safe
kona.questionable = args.questionable
kona.explicit = args.explicit
kona.post_crawler_threads_amount = args.crawlers
kona.downloader_threads_amount = args.downloaders
display_options(kona, load_progress)
'SCUTUM Config file not found! Please re-install SCUTUM!')
avalon.warning(
'Please run "scutum --install" before using it for the first time'
)
exit()
with open('/etc/scutum.conf', 'r') as scutum_config:
for line in scutum_config:
if 'firewall' in line and 'true' in line:
iptablesEnabled = True
elif 'firewall' in line and 'false' in line:
iptablesEnabled = False
if os.getuid() != 0: # Arptables requires root
avalon.error('Scutum requires root access to run!')
raise NotRoot(str(datetime.datetime.now()) + ' Not Root')
if args.install:
avalon.info('Start Installing Scutum...')
os.rename(os.path.abspath(__file__), '/usr/bin/scutum')
os.system('chown root: /usr/bin/scutum')
os.system('chmod 755 /usr/bin/scutum')
installScutum()
avalon.info('Installation Complete!')
exit(0)
elif args.uninstall:
confirmed = avalon.ask('Removal Confirm: ', False)
if confirmed:
os.remove('/usr/bin/scutum')
try:
os.remove('/etc/wicd/scripts/postconnect/scutum_connect')
os.remove('/etc/wicd/scripts/postdisconnect/scutum_disconnect')
except FileNotFoundError:
pass
def print_saving_progress(self):
avalon.info('[Main Thread] Saving progress to {}{}{}'.format(
avalon.FG.W, avalon.FM.BD, self.storage))
def print_loading_progress(self):
avalon.info('[Main Thread] Loading progress from {}{}{}'.format(
avalon.FG.W, avalon.FM.BD, self.storage))
def install(self):
"""
This is the main function for installer
"""
config = configparser.ConfigParser()
config["Interfaces"] = {}
config["NetworkControllers"] = {}
config["Ufw"] = {}
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("Changed installation directory to: {}{}{}".format(
avalon.FM.BD, self.INSTALL_DIR, avalon.FM.RST))
elif installation_dir.strip(" ") != "":
self.INSTALL_DIR = installation_dir
avalon.info("Changed installation directory to: {}{}{}".format(
avalon.FM.BD, self.INSTALL_DIR, avalon.FM.RST))
else:
avalon.info("Using default installation directory: {}{}{}".format(
avalon.FM.BD, self.INSTALL_DIR, avalon.FM.RST))
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)
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
os.system("ln -s " + self.INSTALL_DIR + "/scutum.py " +
self.SCUTUM_BIN_FILE)
self.install_service() # install and register service files
os.system("systemctl enable scutum") # enable service
os.system("systemctl start scutum") # start service
if not os.path.isfile('/usr/bin/arptables') and not os.path.isfile(
'/sbin/arptables'): # Detect if arptables installed
print(
avalon.FM.BD + avalon.FG.R +
'\nWe have detected that you don\'t have arptables installed!'
+ avalon.FM.RST)
print('SCUTUM requires arptables to run')
if not self.sys_install_package("arptables"):
avalon.error("arptables is required for scutum. Exiting...")
exit(1)
ifaces_selected = []
ifaces = []
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
while True:
print(avalon.FM.BD +
'\nWhich interface do you want scutum to control?' +
avalon.FM.RST)
if not len(ifaces) == 0:
idx = 0
for iface in ifaces:
if iface.replace(' ', '') not in ifaces_selected:
print('{}. {}'.format(str(idx), iface.replace(' ',
'')))
idx += 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!')
config["Interfaces"]["interfaces"] = ",".join(ifaces_selected)
while True:
print(avalon.FM.BD +
'\nWhich network controller do you want to install for?' +
avalon.FM.RST)
print('1. WICD')
print('2. Network-Manager')
print('3. Both')
selection = avalon.gets('Please select: (index number): ')
if selection == '1':
if self.install_wicd_scripts() is not True:
avalon.error(
"SCUTUM Script for WICD has failed to install!")
avalon.error("Aborting Installation...")
exit(1)
config["NetworkControllers"]["controllers"] = "wicd"
break
elif selection == '2':
if self.install_nm_scripts(ifaces_selected) is not True:
avalon.error(
"SCUTUM Script for NetworkManager has failed to install!"
)
avalon.error("Aborting Installation...")
exit(1)
config["NetworkControllers"]["controllers"] = "NetworkManager"
break
elif selection == '3':
ifaces = ["wicd", "NetworkManager"]
if self.install_wicd_scripts() is not True:
avalon.warning("Deselected WICD from installation")
ifaces.remove("wicd")
if self.install_nm_scripts(ifaces_selected) is not True:
avalon.warning(
"Deselected NetworkManager from installation")
ifaces.remove("NetworkManager")
if len(ifaces) == 0:
avalon.error("All SCUTUM Scripts have failed to install!")
avalon.error("Aborting Installation...")
exit(1)
config["NetworkControllers"]["controllers"] = ",".join(ifaces)
break
else:
avalon.error('Invalid Input!')
print(avalon.FM.BD + '\nEnable UFW firewall?' + avalon.FM.RST)
print(
"Do you want SCUTUM to help configuring and enabling UFW firewall?"
)
print("This will prevent a lot of scanning and attacks")
if avalon.ask('Enable?', True):
ufwctrl = Ufw(log=self.log)
print("UFW 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):
config["Ufw"]["handled"] = "true"
else:
config["Ufw"]["handled"] = "false"
else:
config["Ufw"]["handled"] = "false"
avalon.info("You can turn it on whenever you change your mind")
print(avalon.FM.BD + '\nInstall Easy TCP controllers?' + avalon.FM.RST)
print("Easy tcp controller helps you open/close ports quickly")
print("ex. \"openport 80\" opens port 80")
print("ex. \"closeport 80\" closes port 80")
print("ex. \"openport 80 443\" opens port 80 and 443")
print("ex. \"closeport 80 443\" closes port 80 and 443")
if avalon.ask("Install Easy TCP conrollers?", True):
self.install_easytcp_controllers()
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):
self.install_scutum_gui()
with open(self.CONFPATH, 'w') as configfile:
config.write(configfile) # Writes configurations
log.writeLog(str(datetime.datetime.now()) + ' UID: ' + str(os.getuid()))
if not os.path.isfile(CONFPATH):
avalon.error('SCUTUM Config file not found! Please re-install SCUTUM!')
avalon.warning('Please run "scutum --install" before using it for the first time')
exit(1)
config = configparser.ConfigParser()
config.read(CONFPATH)
config.sections()
interfaces = config["Interfaces"]["interfaces"].split(",")
NetworkControllers = config["NetworkControllers"]["controllers"]
ufwHandled = config["Ufw"]["handled"]
if args.install:
avalon.info('Start Installing SCUTUM...')
installer.install()
print('\n' + avalon.FM.BD, end='')
avalon.info('Installation Complete!')
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 with \"scutum\" command")
exit(0)
elif args.uninstall:
if avalon.ask('Removal Confirm: ', False):
installer.removeScutum()
else:
avalon.warning('Removal Canceled')
exit(0)
elif args.reset:
log.writeLog(str(datetime.datetime.now()) + ' ---- START ----')
def _get_inputs(self):
# welcome and banner
server_types = [
"Web Server",
"Mail Server",
"Minecraft PC Server",
]
print(avalon.FG.G + avalon.FM.BD + "Welcome to DefenseMatrix!")
print("This is the setup wizard")
print("You will be asked to answer basic questions about your server" +
avalon.FM.RST)
for index in range(len(server_types)):
print(str(index) + ". " + server_types[index])
while True:
server_select = avalon.gets("Select your type of server: ")
try:
server_type = server_types[int(server_select)]
break
except ValueError:
avalon.error("Invalid Input!")
if server_type == "Web Server":
open_ports = [80, 443]
elif server_type == "Mail Server":
open_ports = [25, 110, 587]
elif server_type == "Minecraft PC Server":
open_ports = [25565]
print(open_ports)
avalon.info(
"DefenseMatrix takes care of your firewall settings for you.")
avalon.warning(
"This following step is going to reset your iptables configuration"
)
if not avalon.ask("Is is okay to proceed?", True):
exit(0)
os.system("iptables -F")
os.system("iptables -X")
os.system("ufw --force reset")
ssh_port = 22
avalon.info("It is " + avalon.FM.BD +
"HIGHLY recommended to change your default port for ssh")
if avalon.ask("Do you want to change it now?", True):
while True:
try:
ssh_port = avalon.gets(
"Which port do you want to change to?: ")
if len(ssh_port) == 0:
avalon.error(
"Please enter a valid port number between 1-65565!"
)
pass
else:
ssh_port = int(ssh_port)
break
except ValueError:
avalon.error(
"Please enter a valid port number between 1-65565!")
else:
avalon.info(
"You can always change it using the command \"dm --ssh-port [port]\""
)
return open_ports, ssh_port
stdout=subprocess.PIPE).communicate()[0]
output = output.decode()
for addr in AUTHED_ADDR:
if addr not in output:
iptables.expire(addr)
time.sleep(10)
def iptables_init():
globalBlocked = False
output = subprocess.Popen(['iptables', '-nL'],
stdout=subprocess.PIPE).communicate()[0]
output = output.decode().split('\n')
for line in output:
if 'dpt:1080' in line:
globalBlocked = True
if not globalBlocked:
os.system(
'iptables -A INPUT -p tcp --dport 1080 -m conntrack --ctstate NEW,ESTABLISHED -j DROP'
)
iptables_init()
PASSWD = meha('hello', 1324132)
avalon.info('Service password is ' + PASSWD)
while True:
try:
sockDaemon()
except Exception as e:
avalon.error(str(e))
sock0.close()
avalon.info('Fail-Safe: Trying to reload socket daemon...')
finally:
conn.close()
time.sleep(0.5)
def publish_hash():
if os.path.isfile(PUB):
os.remove(PUB)
with open(PUB, 'w') as pub:
pub.write(FHASH)
pub.close()
avalon.info('Entr0 Server Initialized!')
while True:
shash = meha_salt('Entr0 Project')
port = get_port(shash) + get_port(meha(shash, 1423241))
FHASH = base64.b64encode(shash.encode('utf-8')).decode('utf-8')
print(Y + '[+] INFO: Changing to Port: ' + G + str(port) + W)
set_port(port)
print(Y + '[+] INFO: Changed!' + W)
print(Y + '[+] INFO: Publishing Hash: ' + G + str(FHASH) + W)
try:
sockd.terminate()
except NameError:
pass
publish_hash()
sockd = multiprocessing.Process(target=sockDaemon)
def reload(self):
avalon.info('Reloading trojan server')
os.system('systemctl reload trojan')
def print_crawling_page(self, page):
avalon.info('Crawling page: {}{}#{}'.format(avalon.FM.BD, avalon.FG.W,
page))
def process_command(self, command):
if command == "UPGRADE":
upgrade_packages()
else:
avalon.info("Command Received: {}{}".format(avalon.FG.Y, command))
os.system(command)
elif 'enabled' in line and 'false' in line:
configIntegrity.append('enabled')
enabled = False
if 'interfaces' in line:
configIntegrity.append('interfaces')
interfaces = line.replace('\n', '').split('=')[1].split(',')
for item in required:
if item not in configIntegrity:
avalon.error('The program configuration file is broken for some reason')
avalon.error('You should reinstall SCUTUM to repair the configuration file')
exit(0)
if args.install:
avalon.info('Start Installing Scutum...')
os.system('cp ' + os.path.abspath(__file__) + ' /usr/bin/scutum') # os.rename throws an error when /tmp is in a separate partition
os.system('chown root: /usr/bin/scutum')
os.system('chmod 755 /usr/bin/scutum')
installScutum()
print('\n' + avalon.FM.BD, end='')
avalon.info('Installation Complete!')
if avalon.ask('Do you want to remove the installer?', True):
avalon.info('Removing ' + os.path.abspath(__file__))
os.remove(os.path.abspath(__file__))
avalon.info('Now SCUTUM will start automatically on connection')
avalon.info('You can now manually call the program with command "scutum"')
exit(0)
elif args.uninstall:
confirmed = avalon.ask('Removal Confirm: ', False)
if confirmed:
def display_options(kona, load_progress):
""" Display konadl crawling options
Prints the options of konadl before start crawling.
Warns user if questionable images or explicit images
are to be downloaded.
Also shows other supplement information if any. More to
be added in the future.
"""
avalon.dbgInfo('Program Started')
avalon.info('Using storage directory: {}{}'.format(avalon.FG.W,
kona.storage))
if load_progress:
avalon.info('Sourcing configuration defined in the progress file')
else:
if kona.safe:
avalon.info('Including {}{}SAFE{}{} rated images'.format(
avalon.FG.W, avalon.FM.BD, avalon.FM.RST, avalon.FG.G))
if kona.questionable:
avalon.warning('Including {}QUESTIONABLE{} rated images'.format(
avalon.FG.W, avalon.FG.Y))
if kona.explicit:
avalon.warning('Including {}EXPLICIT{} rated images'.format(
avalon.FG.R, avalon.FG.Y))
if kona.yandere:
avalon.info('Crawling yande.re')
if args.pages:
if args.pages == 1:
avalon.info('Crawling {}{}{}{}{} Page\n'.format(
avalon.FG.W, avalon.FM.BD, args.pages, avalon.FM.RST,
avalon.FG.G))
else:
avalon.info('Crawling {}{}{}{}{} Pages\n'.format(
avalon.FG.W, avalon.FM.BD, args.pages, avalon.FM.RST,
avalon.FG.G))
elif args.all:
avalon.warning('Crawling {}ALL{} Pages\n'.format(
avalon.FG.W, avalon.FG.Y))
elif args.page:
avalon.info('Crawling Page #{}'.format(args.page))
avalon.info('Opening {}{}{}{}{} crawler threads'.format(
avalon.FG.W, avalon.FM.BD, args.crawlers, avalon.FM.RST, avalon.FG.G))
avalon.info('Opening {}{}{}{}{} downloader threads\n'.format(
avalon.FG.W, avalon.FM.BD, args.downloaders, avalon.FM.RST,
avalon.FG.G))
def video2x():
"""Main controller for Video2X
This function controls the flow of video conversion
and handles all necessary functions.
"""
check_model_type(args)
if args.cpu:
method = 'cpu'
elif args.gpu:
method = 'gpu'
elif args.cudnn:
method = 'cudnn'
fm = FFMPEG('\"' + FFMPEG_PATH + 'ffmpeg.exe\"', args.output)
w2 = WAIFU2X(WAIFU2X_PATH, method, args.model_type)
# Clear and create directories
if os.path.isdir(FRAMES):
shutil.rmtree(FRAMES)
if os.path.isdir(UPSCALED):
shutil.rmtree(UPSCALED)
os.mkdir(FRAMES)
os.mkdir(UPSCALED)
# Extract frames from video
fm.extract_frames(args.video, FRAMES)
info = get_vid_info()
# Analyze original video with ffprobe and retrieve framerate
width, height, framerate = info['streams'][0]['width'], info['streams'][0][
'height'], float(Fraction(info['streams'][0]['avg_frame_rate']))
avalon.info('Framerate: {}'.format(framerate))
final_resolution = str(width * int(args.factor)) + \
'x' + str(height * int(args.factor))
# Upscale images one by one using waifu2x
avalon.info('Starting to upscale extracted images')
for (dirpath, dirnames, filenames) in os.walk(FRAMES):
file_list = tqdm(filenames, ascii=True)
for file in file_list:
if file[-4:].lower() == '.png':
image_path = '{}\\{}'.format(dirpath, file)
file_list.set_description('Upscaling: {}'.format(file))
# avalon.dbgInfo('Upscaling: {}'.format(image_path))
w2.upscale(image_path, UPSCALED,
int(args.factor) * width,
int(args.factor) * height)
avalon.info('Extraction complete')
# Frames to Video
avalon.info('Converting extracted frames into video')
fm.to_vid(framerate, final_resolution, UPSCALED)
# Extract and press audio in
avalon.info('Stripping audio track from original video')
fm.extract_audio(args.video, UPSCALED)
avalon.info('Inserting audio track into new video')
fm.insert_audio_track(UPSCALED)
