Allows flag updating using SSH (For Linux) and SMB (For Windows). Authenticated, and allows the use of password/private key over ssh with user specific and port specific authentication. Allows flag_format to be changed at runtime by modifying the configuration file.

flag_updator.py

# For enums
from enum import Enum

# For parsing and handling
import json
import sys
import random
import string
import logging
import tempfile
from base64 import b64decode

# For connection management
import paramiko
from smb.SMBConnection import SMBConnection

class Config:
    """Config class for storing configurations.

    Attributes:
        format: Flag format to be used for flag generation. Default is "CTF"
        length: Length of the flag to be generated. Default is 32
        log_level: Log level for the logger. Default is logging.DEBUG

    """
    format = "CTF"
    length = 32
    log_level = logging.DEBUG

    def change_log_level(level : int) -> None:
        """Changes the log level for the logger.

        Args:
            level: Log level to be set for the logger.
        """
        Config.log_level = level
        logger.setLevel(Config.log_level)


class Logger(object):

    """Logger class for logging.

    Attributes:
        name: Name of the logger
        level: Level of the logger

    Methods:
        get_logger: Returns a logger object
    
    Classes:
        Formatter: Logging Formatter to add colors and count warning / errors
    
    """

    class Formatter(logging.Formatter):
            
            """Logging Formatter to add colors and count warning / errors"""
            grey = "\x1b[38;20m"
            yellow = "\x1b[33;20m"
            red = "\x1b[31;20m"
            green = "\x1b[32;20m"
            bold_red = "\x1b[31;1m"
            reset = "\x1b[0m"
            format = "%(asctime)s %(levelname)s %(message)s"

            FORMATS = {
                logging.DEBUG: grey + format + reset,
                logging.INFO: green + format + reset,
                logging.WARNING: yellow + format + reset,
                logging.ERROR: red + format + reset,
                logging.CRITICAL: bold_red + format + reset
            }
    
            def format(self, record : logging.LogRecord) -> str:

                """Format the log record.

                Args:
                    record: Log record to be formatted
                
                Returns:
                    Formatted log record
                """
                log_fmt = self.FORMATS.get(record.levelno)
                formatter = logging.Formatter(log_fmt)
                return formatter.format(record)
    
    @staticmethod
    def get_logger(name: str, level: int = logging.DEBUG) -> logging.Logger:
        """Returns a logger object.

        Args:
            name: Name of the logger
            level: Level of the logger

        Returns:
            A logger object.

        """
        logger = logging.getLogger(name)
        logger.setLevel(level)
        handler = logging.StreamHandler(sys.stdout)
        handler.setLevel(level)
        formatter = Logger.Formatter()
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        return logger

logger = Logger.get_logger(__name__, level=Config.log_level)

class ENUMS(object):
    
    """ENUMS class for storing enums.

    Attributes:
        OS: OS types
        FlagType: Flag types

    """
    class OS(Enum):
        """Enum for OS types.

        Attributes:
            WINDOWS: Windows OS
            LINUX: Linux OS
            MAC: Mac OS
            UNKNOWN: Unknown OS
        """
        WINDOWS = "windows"
        LINUX = "linux"
        UNKNOWN = "unknown"

    class FlagType(Enum):
        """Enum for flag types.

        Attributes:
            MACHINE: Machine specific flag (That is generated using key-machine attributes [Honestly, Idk why...])
            GENERIC: Generic flag that can be used for different purposes
        """
        MACHINE = 1
        GENERIC = 2

    class ConnectionType(Enum):

        """Enum for connection types.

        Attributes:
            SSH: SSH connection
            SMB: SMB connection
        """
        SSH = "ssh"
        SMB = "smb"

    class FlagType(Enum):

        """Enum for flag types.

        Attributes:
            USER: User flag
            ROOT: Root flag
        """

        USER = "user"
        ROOT = "root"

class BaseManager(object):

    FLAG_REPLACE_ATTR = "[[FLAG]]"
    FILE_REPLACE_ATTR = "[[FILE]]"
    CONT_REPLACE_ATTR = "[[CONT]]"

    @staticmethod
    def __generate_random_string(
        length: int
    ) -> str:
        """Generates a random string of given length.

        Args:
            length: Length of the random string to be generated.

        Returns:
            A random string of given length.

        """
        return "".join(random.choice(string.ascii_letters + string.digits) for _ in range(length))
    
    @staticmethod
    def generate_flag(
        flag_format: str = None,
        flag_len: int = None
    ) -> str:
        
        """Generates a flag based on the given flag format.

        Args:
            flag_format: Flag format to be used for flag generation. Default is "CTF"
            flag_len: Length of the flag to be generated. Default is 32
            
        Returns:
            A flag based on the given flag format.
        """

        """ A very bad of testing, but it works..."""
        if flag_format is None:
            flag_format = Config.format

        if flag_len is None:
            flag_len = Config.length

        return flag_format + "{" + BaseManager.__generate_random_string(flag_len) + "}"

class SSHHandler(BaseManager):

    """SSHHandler class for handling SSH connections.

    Attributes:
        ip_address: IP address of the target machine
        username: Username for the SSH connection
        password: Password for the SSH connection
        priv_key: Private key for the SSH connection
        port: Port for the SSH connection
        os: OS type of the target machine
        is_docker: Whether the machine is a docker container or not
        cmd: A dictionary of commands to be executed on the target machine

    Raises:
        Exception: Cannot have both password and private key

    """

    BASE_CMD   = f"echo {BaseManager.FLAG_REPLACE_ATTR} > {BaseManager.FILE_REPLACE_ATTR} && echo $?"
    DOCKER_CMD = f"docker exec -i {BaseManager.CONT_REPLACE_ATTR} sh -c '{BASE_CMD}'"

    def __init__(
        self,
        ip_address: str,
        username: str,
        password: str,
        priv_key: str,
        port: int = 22
    ):
        self.ip_address = ip_address
        self.username = username
        self.password = password
        self.priv_key = priv_key
        self.port = port

        if self.priv_key != "" and self.password != "":
            raise Exception("Cannot have both password and private key")
    
    def session(self) -> paramiko.SSHClient:
        """Validates the SSH connection.

        Returns:
            True if the SSH connection is valid, False otherwise.
        """
        try:
            logger.debug(f"Obtaining SSH connection to {self.ip_address}:{self.port}")
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            if self.priv_key != "":
                """Storing the private key in a temporary file and using it for the SSH connection"""
                with tempfile.NamedTemporaryFile() as tmp:
                    tmp.write(b64decode(self.priv_key))
                    tmp.flush()
                    tmp.seek(0)
                    self.key_file = tmp.name
                    logger.debug(f"Private key stored in {self.key_file}")

                    ssh.connect(
                        self.ip_address,
                        port=self.port,
                        username=self.username,
                        key_filename=self.key_file
                    )
            else:
                ssh.connect(
                    self.ip_address,
                    port=self.port,
                    username=self.username,
                    password=self.password
                )
            return ssh
        except Exception as e:
            logger.error(f"Error while obtaining an SSH connection: {e}")
            exit(1)

    def update_flag(
        self,
        flag: str,
        file: str,
        is_docker: bool = False,
        container_name: str = ""
    ) -> bool:
        
        """Updates the flag on the target machine.

        Args:
            flag: Flag to be updated on the target machine
            file: File to be updated on the target machine
            is_docker: Whether the machine is a docker container or not
            container_id: Container ID of the docker container

        Returns:
            True if the flag is updated successfully, False otherwise.
        """

        s = self.session()
        if not s:
            return False
        
        if is_docker and container_name == "":
            logger.error("Container Name cannot be empty for docker containers")
            return False

        if is_docker:
            cmd = self.DOCKER_CMD
        else:
            cmd = self.BASE_CMD

        cmd = cmd.replace(BaseManager.FLAG_REPLACE_ATTR, flag)
        cmd = cmd.replace(BaseManager.FILE_REPLACE_ATTR, file)

        if is_docker:
            cmd = cmd.replace(BaseManager.CONT_REPLACE_ATTR, self.container_id)

        logger.debug(f"Executing command: {cmd}")
        stdin, stdout, stderr = s.exec_command(cmd)

        if stdout.channel.recv_exit_status() != 0:
            logger.error(f"Error while executing command: {stderr.read().decode()}")
            return False
        
        logger.info(f"Flag updated successfully: {flag}")
        return True

class SMBHandler(BaseManager):
    
    """SMBHandler class for handling SMB connections.

    Attributes:
        ip_address: IP address of the target machine
        username: Username for the SMB connection
        password: Password for the SMB connection
        port: Port for the SMB connection
        os: OS type of the target machine
        share_name: Share name for the SMB connection
    """

    def __init__(
        self,
        ip_address: str,
        username: str,
        password: str,
        port: int = 445,
        share_name: str = "C$"
    ):
        self.ip_address = ip_address
        self.username = username
        self.password = password
        self.port = port
        self.share_name = share_name

    def __is_share__(
        self,
        share_name : str
    ):
        for share in self.shares:
            if share_name.lower() == share.name.lower():
                return True
            
        return False
    
    def session(self) -> SMBConnection:
        """Validates the SMB connection.

        Returns:
            True if the SMB connection is valid, False otherwise.
        """
        try:
            logger.debug(f"Obtaining SMB connection to {self.ip_address}:{self.port}")
            conn = SMBConnection(
                self.username,
                self.password,
                my_name="",
                remote_name=self.ip_address,
                use_ntlm_v2=True,
                is_direct_tcp=True
            )
            conn.connect(self.ip_address, self.port)
            self.shares = conn.listShares()
            return conn
        except Exception as e:
            logger.error(f"Error while obtaining an SMB connection: {e}")
            exit(1)

    def update_flag(
        self,
        flag: str,
        file: str
    ) -> bool:
        
        """Updates the flag on the target machine.

        Args:
            flag: Flag to be updated on the target machine
            file: File to be updated on the target machine

        Returns:
            True if the flag is updated successfully, False otherwise.
        """

        s = self.session()
        if not s:
            return False
        
        if not self.__is_share__(self.share_name):
            logger.error(f"Share {self.share_name} not found on {self.ip_address}")
            exit(1)

        logger.debug(f"Updating flag on {self.ip_address}:{self.port}/{self.share_name}/{file}")
        try:
            with tempfile.NamedTemporaryFile() as tmp:
                tmp.write(flag.encode())
                tmp.flush()
                tmp.seek(0)
                s.storeFile(
                    self.share_name,
                    file,
                    tmp
                )
        except Exception as e:
            logger.error(f"Error while updating flag: {e}")
            return False
        
        logger.info(f"Flag updated successfully: {flag}")
        return True

class FlagUpdator(BaseManager):

    """FlagUpdator class for updating flags.

    Attributes:
        machine_file: Path to the machines.json file
        machines: A dictionary of machines
    """

    machine_file = ""
    machines = {}

    class MachineInfo(object):

        """MachineInfo class for storing machine information.

        Attributes:
            name: Name of the machine
            ip: IP address of the machine
            port: Port of the machine
            os: OS type of the machine
            user: User information of the machine [port, docker, path]
            root: Root information of the machine [port, docker, path]
            username: Username for the machine
            password: Password for the machine
            connType: Connection type for the machine
            
        """

        def __init__(
            self,
            name : str = "",
            ip : str = "",
            port : int = 22,
            os : ENUMS.OS = ENUMS.OS.UNKNOWN,
            is_docker : bool = False,
            user : dict = {},
            root : dict = {},
            username : str = "",
            password : str = "",
            priv_key : str = "",
            connType : ENUMS.ConnectionType = ENUMS.ConnectionType.SSH,
            share_name : str = "C$"
        ):
            self.name = name
            self.ip = ip
            self.port = port
            self.os = os
            self.is_docker = is_docker
            self.user = user
            self.root = root
            self.username = username
            self.password = password
            self.connType = connType
            if self.connType == ENUMS.ConnectionType.SSH:
                self.priv_key = priv_key
            self.share_name = share_name

        def __str__(self):
            rt = f"MachineInfo(name={self.name}, ip={self.ip}, port={self.port}, os={self.os}, user={self.user}, root={self.root}, username={self.username}, PWD, connection_type={self.connType}, SHARE_NAME)"

            if self.connType == ENUMS.ConnectionType.SSH:
                if self.priv_key != "":
                    rt = rt.replace("PWD", f"priv_key={self.priv_key}")
                else:
                    rt = rt.replace("PWD", f"password={self.password}")
                rt = rt.replace(", SHARE_NAME", "")

            elif self.connType == ENUMS.ConnectionType.SMB:
                rt = rt.replace("PWD", f"password={self.password}")
                rt = rt.replace("SHARE_NAME", f"share_name={self.share_name}")
        
        def __repr__(self):
            return self.__str__()

    def __init__(
        self,
        json_file="machines.json"
    ):
        try:
            self.machines = json.load(open(json_file, "r"))["Machines"]
        except Exception as e:
            logger.error(f"Error while loading {json_file}: {e}")
            sys.exit(1)
        self.machine_file = json_file
        logger.debug(f"Loaded {len(self.machines)} machines from {json_file}")

    def __parse__(
        self,
        machine_name: str
    ) -> MachineInfo:
        if machine_name not in self.machines:
            logger.error(f"Machine {machine_name} not found in {self.machine_file}")
            exit(1)
        
        machine = self.machines[machine_name]

        if machine["os"].lower() != ENUMS.OS.LINUX.value and machine["os"].lower() != ENUMS.OS.WINDOWS.value:
            logger.error(f"Unsupported OS: {machine['os']}")
            exit(1)

        # Parse to machineINfo:
        info = FlagUpdator.MachineInfo()
        info.name = machine_name
        info.ip = machine["ip"]
        info.port = machine["port"]
        info.os = ENUMS.OS(machine["os"].lower())
        info.user = machine["user"]
        info.root = machine["root"]
        info.username = machine["username"]

        try:
            info.connType = ENUMS.ConnectionType(machine["type"].lower())
        except Exception as e:
            logger.error(f"Invalid connection type for {machine_name}: {e}")
            exit(1)

        try:
            info.password = machine["password"]
        except KeyError:
            if info.connType == ENUMS.ConnectionType.SSH:
                priv = machine.get("priv_key", "")
                if priv == "":
                    logger.error(f"Invalid password or private key for {machine_name}")
                    exit(1)
                else:
                    info.priv_key = priv
                    info.password = ""

        info.share_name = machine.get("share_name", "C$")

        if info.connType == ENUMS.ConnectionType.SMB:
            if info.password == "":
                logger.error(f"Invalid password for {machine_name}")
                exit(1)
            
            if info.priv_key != "":
                logger.error(f"Cannot have private key for {machine_name} as the connection is through SMB")
                exit(1)

        if type(info.user) is not dict:
            logger.error(f"Invalid user attribute for {machine_name}")
            exit(1)

        if type(info.root) is not dict:
            logger.error(f"Invalid root attribute for {machine_name}")
            exit(1)

        if (info.username == "" or info.password == "") and info.priv_key == "":
            logger.error(f"Invalid username or password for {machine_name}")
            exit(1)

        if info.connType == ENUMS.ConnectionType.SSH:
            if info.password != "" and info.priv_key != "":
                logger.error(f"Cannot have both password and private key for {machine_name}")
                exit(1)

        if "port" not in info.user.keys():
            info.user["port"] = info.port
        
        if "port" not in info.root.keys():
            info.root["port"] = info.port

        return info

    def __update__(
        self,
        flag : str,
        info : MachineInfo,
        flag_type : ENUMS.FlagType
    ):
        if info.connType == ENUMS.ConnectionType.SSH:
            ssh = SSHHandler(
                ip_address=info.ip,
                username=info.username,
                password=info.password,
                priv_key=info.priv_key
            )

            def _update(attr):
                if attr["docker"]:
                    cont_name = attr.get("name", "")
                    if cont_name == "":
                        logger.error(f"Invalid container name for {info.name}")
                        exit(1)
                    return ssh.update_flag(
                        flag=flag,
                        file=attr["path"],
                        is_docker=True,
                        container_name=attr["name"]
                    )
                else:
                    return ssh.update_flag(
                        flag=flag,
                        file=attr["path"]
                    )

            if flag_type == ENUMS.FlagType.USER:
                return _update(attr=info.user)
                
            elif flag_type == ENUMS.FlagType.ROOT:
                return _update(attr=info.root)

        elif info.connType == ENUMS.ConnectionType.SMB:
            smb = SMBHandler(
                ip_address=info.ip,
                username=info.username,
                password=info.password,
                share_name=info.share_name
            )

            def _update(attr):
                return smb.update_flag(
                    flag=flag,
                    file=attr["path"]
                )
            
            if flag_type == ENUMS.FlagType.USER:
                return _update(attr=info.user)
            
            elif flag_type == ENUMS.FlagType.ROOT:
                return _update(attr=info.root)

    def __middleware__(
        self,
        machine_name : str,
        type : ENUMS.FlagType,
        flag : str = "",
        flag_format: str = Config.format,
        flag_len: int = Config.length
    ):
        
        info = self.__parse__(machine_name=machine_name)

        if flag_format is None:
            flag_format = Config.format
        
        if flag_len is None:
            flag_len = Config.length

        if not flag:
            flag = BaseManager.generate_flag(flag_format=flag_format, flag_len=flag_len)
            logger.debug(f"Generated flag for {info.name}: {flag}")
        logger.info(f"Adding {type.name} flag for machine {info.name}: {flag}")
        status = self.__update__(flag=flag, info=info, flag_type=type)

        return {
                "flag" : flag,
                "status" : "success" if status else "error"
            }

    def update_user_flag(
        self,
        machine_name: str,
        flag : str = "",
        flag_format: str = None,
        flag_len: int = None
    ):
        return self.__middleware__(machine_name=machine_name, type=ENUMS.FlagType.USER, flag=flag, flag_format=flag_format, flag_len=flag_len)   
    
    def update_root_flag(
        self,
        machine_name: str,
        flag : str = "",
        flag_format: str = None,
        flag_len: int = None
    ):
        return self.__middleware__(machine_name=machine_name, type=ENUMS.FlagType.ROOT, flag=flag, flag_format=flag_format, flag_len=flag_len)

    def add_flags(
        self,
        machine_name: str,
        user_flag = "",
        root_flag = "",
        flag_format: str = None,
        flag_len: int = None,
    ) -> bool:
        
        if flag_format is None:
            flag_format = Config.format

        if flag_len is None:
            flag_len = Config.length
        
        """Adds a flag to the given machine.

        Args:
            machine_name: Name of the machine
            flag_format: Flag format to be used for flag generation. Default is "CTF"
            flag_len: Length of the flag to be generated. Default is 32
            
        Returns:
            True if the flag is added successfully, False otherwise.
        """
        if not user_flag:
            user_flag = BaseManager.generate_flag(flag_format=flag_format, flag_len=flag_len)

        if not root_flag:
            root_flag = BaseManager.generate_flag(flag_format=flag_format, flag_len=flag_len)

        user = self.update_user_flag(
            machine_name=machine_name,
            flag=user_flag
        )

        root = self.update_root_flag(
            machine_name=machine_name,
            flag=root_flag
        )

        return {"user": user, "root": root}

machines.json

{
    "Machines" : {
        "base-img" : {
            "ip" : "192.168.163.168",
            "port" : 22,
            "username" : "root",
            "priv_key" : "LS0tLS1CRUdJTiBPUEVOU1NIIFBSSVZBVEUgS0VZLS0tLS0KYjNCbGJuTnphQzFyWlhrdGRqRUFBQUFBQkc1dmJtVUFBQUFFYm05dVpRQUFBQUFBQUFBQkFBQUJsd0FBQUFkemMyZ3RjbgpOaEFBQUFBd0VBQVFBQUFZRUFvY1hqMG9iaUVFL3JtaDAzdE1icm5mZ2czS2E2Z2JhSGRjczl1bFZpUXJmbTgzbmtZeUo0CjN3anhhb0NBT1YxL0svL25UcjhCcXd1dUhOOFp0Q1cxM2pmeXdMa3VaVWtDNmZ0ZHZYRzRnNGtvS3dJRityOXhWRzJ3aCsKN1ZzdUx4WVY5bXU0YXpwTFA3WjBBa1h2RjVyMDRNdjUraC8zQVNpaExKeGVZbytnT09jdXNrZ1BBbWRVRzd1NjR1MEhLZApTM21SZ1NMcVVDUHk0dmU0YlROcTJDZ2VvUTVISDZNRWNrNGJrZGUvRGV3Y2tYNklYL1BiSjVGbjdhdndvT2tYRmpueVlwCkJRYnFLUnhvWFlPUUZRazhPY01PbjY1NE9vVFhXMTlVcHloSmdkYzhHeFFFWkdwWlJYRnoxeGhNYlNsTU1UNFVQb1NpcVMKWFc5c3hlbnVNVVFpenNDdFI1ZVl3b3FkUUljRVFyV0ppYUkzSFVZWFNxazB1WHRGSkJqMG9wTUoxMGNtZEdUNUV4N000dgo0bkwvYmtuYVBSMnFQZElRd2xZU0gwUVJHWUUzMFd6ZGJYU25PTnR3RzJEamkxZTh6eEVTeDBTYmtIeWdxVXZzeGt2ZHBjCjE4emV4bUR6U0FBYzBDY3RPNndEelFCbHRmOWpiZkR3YVRjZmh4Z1pBQUFGaVB5Q2p4bjhnbzhaQUFBQUIzTnphQzF5YzIKRUFBQUdCQUtIRjQ5S0c0aEJQNjVvZE43VEc2NTM0SU55bXVvRzJoM1hMUGJwVllrSzM1dk41NUdNaWVOOEk4V3FBZ0RsZApmeXYvNTA2L0Fhc0xyaHpmR2JRbHRkNDM4c0M1TG1WSkF1bjdYYjF4dUlPSktDc0NCZnEvY1ZSdHNJZnUxYkxpOFdGZlpyCnVHczZTeisyZEFKRjd4ZWE5T0RMK2ZvZjl3RW9vU3ljWG1LUG9Eam5MckpJRHdKblZCdTd1dUx0QnluVXQ1a1lFaTZsQWoKOHVMM3VHMHphdGdvSHFFT1J4K2pCSEpPRzVIWHZ3M3NISkYraUYvejJ5ZVJaKzJyOEtEcEZ4WTU4bUtRVUc2aWtjYUYyRAprQlVKUERuRERwK3VlRHFFMTF0ZlZLY29TWUhYUEJzVUJHUnFXVVZ4YzljWVRHMHBUREUrRkQ2RW9xa2wxdmJNWHA3akZFCklzN0FyVWVYbU1LS25VQ0hCRUsxaVltaU54MUdGMHFwTkxsN1JTUVk5S0tUQ2RkSEpuUmsrUk1lek9MK0p5LzI1SjJqMGQKcWozU0VNSldFaDlFRVJtQk45RnMzVzEwcHpqYmNCdGc0NHRYdk04UkVzZEVtNUI4b0tsTDdNWkwzYVhOZk0zc1pnODBnQQpITkFuTFR1c0E4MEFaYlgvWTIzdzhHazNINGNZR1FBQUFBTUJBQUVBQUFGL0RTNjg1Y3IxazZjbjJZT21QbDBVNThWQ0ZVCnRFWS90TkNhVkNXbC9KOWhpQmFab0crd2pmczlGTzFVOGJwVGNMTGdLZXZ6QTNiVTcwYTNTbEtZdG1WRS9OMUJ0Zkk3VjkKZDNtclg2c0FlSkFsbFJycDJoYi9wWTVNSFNkRWZvb0hnenB4QjRCL2pyR2hsZHhsUjF1VDFrNDJGbmJrL3d2Q09Sci9YZgpxNlRydk52blFzVXczQU1Lb1E4aHUyRW01L3BhN3hKK1NkMkxsMDNGMnA1b3J2dHZ1blNEdGJwTVB4c254ajZBWHhTZTlqCjA3UFowZVF0OVRNM2ZpNEZ3b3BnZFFPTHRFd2NaU1lzU0dSanpSbmpJRW5SRTQrZkJNc1ZMRWhaZ243SlBvTHpkTHlHSjEKcHN1MG0veGVVZk45QWZmTGpPQlpDS1VjdytTdm9BUTBHNmNhVzdyT0JXbVI1aHdaWWNQUnMvbDVSVFU1SXgvNzBHTytobAptcnUxUi9ITzdJdnN6Q3hsU1RnSW9XNmNGNEZrRE5BRWlDeEtnd2NSSU9QSll4dXZoSEkrcVRkWnlJczMxMUhCWDRXb2ZjCi9IZzIxVVJSeThhdVJFNVArWG5VOEkwR0NhbXI3L0FIU1k1bWJRa1pLeE4zaTVlNEFYS0k3NUd5d2NCWjVNM1dFQUFBREEKT1RKc3Y0T3kyMjNYSXpOZEpFWGhiRzRsaUkwM0YyVGZJOFJsa1VMbUJTMzRGSmQ4NWJXQWsvUXJnVGtSTE05QWdXeDF3QQpPaUpEbk9DVVFTU3RETGo4YTkvcXE3eENWaWd3bnp1S2ZaRkpKcUJ3ajFhSXdNbFRHcmxuc1RvaDlZbU5XaG5yVUI5OE43CkNNLzM3OXRpNm8vN3ZKTGEvMmFaWUhWdVhldjFnbHZOU0FkZDNWNFNvcThZcUhmaEpLSlA1Y29QdzdJQkVra05iRWVjOCsKTkhaUUxYOVRNOTlObHhrazZtSXJnWW5BOHlFbit0bnNTSHY1eUxJM2lGSHJnZkFBQUF3UURQL2NMM2FQUWsyWHBONHl6VApGSDNzcExTUWE4UTV3UFlSOFRuRGRudm84cHVsa3FIbU8wMzFZVmhpTjl1RjlwV1RodEpyNFQ0TlJuRXIwaDUxQjF4dk1lCjNSR0M0VUZkNDZXSWlKVmtrazQ4QWV5M1BQRlhTU2pHY1l2VzI4WFJ4NjhkQ3Z6UU13dXcwVGZwQXlhWU1QdFVPMGNEdXAKUTFLTEZYd081U1ptYTdwOUNHRnRXUHRZaVp3R0NzZzhoMjIva2MydXQ5ejJheWhValh4ZWRRTlVJcm9KWDR5UTJWYThkRgpjeS9VMXZDSS9OMGFFTzM3Q2pqMFFjVG50M2pYa0FBQURCQU1jZEZaQmdxaTVKbEUyMk5qTDZQUG4wK2VnM3dZOEpKRzJNClpsV2M0dmIvVlhkM1puYytqT2wyUlNSdGpEOWRCUVB5cXV4MkhWTEFWSnNNOG9NT1JVTHlsZ2tDSUFGN2Q3c3NCQVJIVysKS0dPQUZiRnR6OGUxZ1NHQ0tveDB6NjBkeWxZcmFsUXp1M3pJRGxweVlrQkVDTlpLY0NOaTk5UnI2ZENBKy9penNWZWJIRgpNSUNENkxmMHk4ZFpvS0tsUnFWaXh6ajNKK1JqbWs2NUE1Sjl0SStIdWRHNytFNUw0eGZvTlNvWXFqSi9kZm9xQTN2aEpmCmxES2RkQXNKdFhvUUFBQUJKeWIyOTBRSFZpZFc1MGRTMXpaWEoyWlhJQkFnPT0KLS0tLS1FTkQgT1BFTlNTSCBQUklWQVRFIEtFWS0tLS0tCg==",
            "user" : {
                "docker" : false,
                "name" : "user",
                "path" : "/home/ubuntu/user.txt"
            },
            "root" : {
                "docker" : false,
                "path" : "/root/root.txt",
                "port" : 22
            },
            "os" : "linux",
            "type" : "ssh"
        },
        "base-img-win" : {
            "ip" : "192.168.163.148",
            "port" : 22,
            "username" : "Administrator",
            "password" : "password-goes-here",
            "user" : {
                "docker" : false,
                "name" : "user",
                "path" : "user.txt"
            },
            "root" : {
                "docker" : false,
                "path" : "Users/user/Desktop/root.txt"
            },
            "os" : "windows",
            "type" : "smb",
            "share_name" : "C"
        }
    }
}

requirements.txt

paramiko==3.3.1
paramiko_ng==2.8.10
pysmb==1.2.9.1

Usage:

from flag_updator import FlagUpdator, Config

updator = FlagUpdator(json_file="machines.json")

print("*" * 50)
print(">> I just want to update BOTH flags..")
flags = updator.add_flags(machine_name="base-img-win")

print("*" * 50)
print(">> I just want to update the user flag...")
updator.update_user_flag(machine_name="base-img-win")

print("*" * 50)
print(">> I just want to update the root flag with a new format...")
Config.format = "FLAG_UPDATE"
updator.update_root_flag(machine_name="base-img-win")

Output is something like this:

2023-10-18 02:50:09,315 DEBUG Loaded 7 machines from machines.json
**************************************************
>> I just want to update BOTH flags..
2023-10-18 02:50:09,315 INFO Adding USER flag for machine base-img-win: CTF{ybQjgJ6h5sfpnxJLYoBqxAD0i9dDPXSn}
2023-10-18 02:50:09,315 DEBUG Obtaining SMB connection to 192.168.163.148:445
2023-10-18 02:50:09,323 DEBUG Updating flag on 192.168.163.148:445/C/user.txt
2023-10-18 02:50:09,325 INFO Flag updated successfully: CTF{ybQjgJ6h5sfpnxJLYoBqxAD0i9dDPXSn}
2023-10-18 02:50:09,325 INFO Adding ROOT flag for machine base-img-win: CTF{POfLfeO0bwAdVTf5YuT9MCROknrIGUyu}
2023-10-18 02:50:09,325 DEBUG Obtaining SMB connection to 192.168.163.148:445
2023-10-18 02:50:09,333 DEBUG Updating flag on 192.168.163.148:445/C/Users/user/Desktop/root.txt
2023-10-18 02:50:09,337 INFO Flag updated successfully: CTF{POfLfeO0bwAdVTf5YuT9MCROknrIGUyu}
**************************************************
>> I just want to update the user flag...
2023-10-18 02:50:09,337 DEBUG Generated flag for base-img-win: CTF{eVPk5UBb6565v5ZFZSoW30admoQ0dyHm}
2023-10-18 02:50:09,337 INFO Adding USER flag for machine base-img-win: CTF{eVPk5UBb6565v5ZFZSoW30admoQ0dyHm}
2023-10-18 02:50:09,337 DEBUG Obtaining SMB connection to 192.168.163.148:445
2023-10-18 02:50:09,346 DEBUG Updating flag on 192.168.163.148:445/C/user.txt
2023-10-18 02:50:09,349 INFO Flag updated successfully: CTF{eVPk5UBb6565v5ZFZSoW30admoQ0dyHm}
**************************************************
>> I just want to update the root flag with a new format...
2023-10-18 02:50:09,349 DEBUG Generated flag for base-img-win: FLAG_UPDATE{DV23oDxyHO7bl6GeZcn0hTFFSC54onG9}
2023-10-18 02:50:09,349 INFO Adding ROOT flag for machine base-img-win: FLAG_UPDATE{DV23oDxyHO7bl6GeZcn0hTFFSC54onG9}
2023-10-18 02:50:09,349 DEBUG Obtaining SMB connection to 192.168.163.148:445
2023-10-18 02:50:09,357 DEBUG Updating flag on 192.168.163.148:445/C/Users/user/Desktop/root.txt
2023-10-18 02:50:09,361 INFO Flag updated successfully: FLAG_UPDATE{DV23oDxyHO7bl6GeZcn0hTFFSC54onG9}