bcrypt.ctf/bcrypter/cli/opt.py

from enum import Enum
from dataclass import dataclass

from bcrypter.lib.other import tryget
from bcrypter.exceptions import (
    BadOptNameError, 
    InvalidOptError, 
    DisjointOverlapError,
)

'''
Overlap dataclass is used to check if two Opts are incompatible
ie same name, same shortform, same longform
'''
@dataclass
class Overlap:
    master: Opt # the Opt we measure overlap against
    name: list[Opt]
    shortform: list[Opt]
    longform: list[Opt]

    '''
    Create an empty Overlap for a master
    '''
    @classmethod
    def empty(cls: Overlap, A: Opt) -> Overlap:
        return cls(A, [], [], [])
    
    '''
    Calculate the overlap of two objects inheriting from Opt 
    '''
    @classmethod
    def of(cls: Overlap, A: Opt, B: Opt) -> Overlap:
        # "equal to and not None"
        eqnn = lambda x, y: x is not None and x == y
        return cls(A,
                   [B] if eqnn(A.name, B.name) else [],
                   [B] if eqnn(A.shortform, B.shortform) else [],
                   [B] if eqnn(A.longform, B.longform) else [])
    
    def __add__(self, O: Overlap) -> Overlap:
        if self.master != O.master:
            raise DisjointOverlapError()
        return Overlap(master, 
                       self.name + O.name,
                       self.shortform + O.shortform,
                       self.longform + O.longform)

    def is_empty(self) -> bool:
        return not (self.name or self.shortform or self.longform):

class ParamType(Enum):
    AbstractParam # only used by Param
    AbstractOpt   # only used by Opt
    Arg
    Flag
    Option

class Param:
    _TYPE: ParamType = ParamType.AbstractParam
    def __init__(self, name: str, description: str) -> None:
        self.name = name
        self.description = description

class Opt(Param):
    _TYPE: ParamType = ParamType.AbstractOpt
    __FAIL_MSG = 'Invalid Opt'
    def __init__(self, *args, name: str = '', description: str = '') -> None:
        # ensure non-empty name kwarg provided
        if not name:
            raise BadOptNameError()
        super().__init__(name, description)
        # result.value is a list of error messages (if failed)
        self.result: Result[None|list[str]] = Result.succeed(None)
        self.shortform: str | None = None
        self.longform: str | None = None
        self._parse_args(args)
            
    '''
    Parse provided *args as short/longform
    NOTE: sets self.result as fail with error message list as value 
    '''
    def _parse_args(self, args: list[str]) -> None:
        # store all short/long/unknown forms although multiple arent allowed
        # used to provide better errors (displaying ALL duplicates)
        shortforms = []
        longforms = []
        unknowns = []
        for arg in args:
            if Opt.is_longform(arg):
                longforms.append(arg)
            elif Opt.is_shortform(arg):
                shortforms.append(arg)
            else:
                unknowns.append(arg)
        self.shortform = tryget(shortforms, 0) # default: None
        self.longform = tryget(longforms, 0)   # default: None
        # check if any duplicate arg forms occurred
        errs = []
        if len(shortforms) > 1:
            errs.append(Opt._fail_duplicate_shortsforms(shortforms))
        if len(longforms) > 1:
            errs.append(Opt._fail_duplicate_longforms(longforms))
        if unknowns:
            errs.append(Opt._fail_unknown_form(unknowns))
        # ensure at least one short/longform given
        if not (shortforms or longforms or unknowns):
            raise InvalidOptError()
        # set Opt.result according to derived errors
        if errs:
            if self.result.is_err():
                self.result.value.extend(errs)
            else:
                self.result = Result.fail(Opt.__FAIL_MSG, value=errs)        

    '''
    Opt.is_shortform and Opt.is_longform check whether an
    arg is in the form ie `-v`, or `--verbose` respectively.
    NOTE: Always check longform before shortform, ie anything
    NOTE: with is_longform will also be is_shortform (im lazy) 
    '''
    @staticmethod
    def is_shortform(arg: str) -> bool:
        return arg.startswith('-')

    @staticmethod
    def is_longform(arg: str) -> bool:
        return arg.startswith('--')
    
    '''
    =========================================================
    Methods for formatting+returning error messages as a way
    of logging the outcome of longform/shortform/etc parsing.
    =========================================================
    '''
    @staticmethod
    def _fail_bad_forms(form_type: str, forms: list) -> str:
        forms = ', '.join[f'\"{f}\"' for f in forms]
        return f'Opt \"{self.name}\" got {form_type}: {forms}'
    
    @staticmethod
    def _fail_duplicate_shortforms(args: list[str]) -> str:
        return Opt._fail_bad_forms('duplicate shortforms', args)

    @staticmethod
    def _fail_duplicate_longforms(args: list[str]) -> str:
        return Opt._fail_bad_forms('duplicate longforms', args)
    
    @staticmethod
    def _fail_unknown_forms(args: [str]) -> str:
        form_type = 'unknown arg'
        if len(args) > 1:
            form_type += 's' # make plural
        return Opt._fail_bad_forms(form_type, args)

    # @staticmethod
    # def _fail_overlap(form_type: str, form: str, name1: str, name2: str) -> str:
    #     return f'Duplicate {form_type} \"{form}\" for Opts \"{name1}\" & \"{name2}\"'

class Arg(Param):
    _TYPE: ParamType = ParamType.Arg
    def __init__(self, name: str = '', description: str = '') -> None:
        super().__init__(name=name, description=description)

class Flag(Opt):
    _TYPE: ParamType = ParamType.Flag
    def __init__(self, *args, name: str = '', description: str = '') -> None:
        super().__init__(name=name, description=description)

class Option(Opt):
    _TYPE: ParamType = ParamType.Option
    def __init__(self, *args, name: str = '', description: str = '') -> None:
        super().__init__(name=name, description=description)