Newer
Older
Hardware / FaultInjection / prereqs / FaultyCat / Modules / CmdInterface.py
0xRoM on 11 Feb 2 KB initial commit
  1. import cmd
  2. import typer
  3. from rich.console import Console
  4. from rich.table import Table
  5.  
  6.  
  7. def is_valid_number(number):
  8. if number < 0:
  9. raise typer.BadParameter("Number must be positive.")
  10. return number
  11.  
  12. class CMDInterface(cmd.Cmd):
  13. intro = "Type help or ? to list commands.\n"
  14. prompt = "?> "
  15. file = None
  16. doc_header = "Commands"
  17. misc_header = "Misc Commands"
  18. undoc_header = "Undocumented Commands"
  19.  
  20. def __init__(self, faulty_worker):
  21. super().__init__()
  22. self.faulty_worker = faulty_worker
  23.  
  24. def do_config(self, args):
  25. """Configure the FaultyCat."""
  26. print("Configuring the FaultyCat...")
  27. table_config = Table(title="Board configuration")
  28. table_config.add_column("Parameter", style="cyan")
  29. table_config.add_column("Value", style="magenta")
  30. table_config.add_row(
  31. "Serial port", f"{self.faulty_worker.board_uart.serial_worker.port}"
  32. )
  33. table_config.add_row(
  34. "Pulse time",
  35. f"{self.faulty_worker.board_configurator.BOARD_CONFIG['pulse_time']}",
  36. )
  37. table_config.add_row(
  38. "Pulse power",
  39. f"{self.faulty_worker.board_configurator.BOARD_CONFIG['pulse_power']}",
  40. )
  41. table_config.add_row("Pulse count", f"{self.faulty_worker.pulse_count}")
  42. Console().print(table_config)
  43.  
  44. def do_set(self, args):
  45. """Set a parameter."""
  46. print("Setting a parameter...")
  47. args_list = args.split()
  48. if args == "help" or args == "?":
  49. print("Available parameters:")
  50. print("\t[time] pulse_time")
  51. print("\t[count] pulse_count")
  52. print("\tport")
  53. return
  54.  
  55. if len(args_list) != 2:
  56. print("Invalid number of arguments.")
  57. return
  58.  
  59. if args_list[0] == "pulse_time" or args_list[0] == "time":
  60. self.faulty_worker.set_pulse_time(is_valid_number(float(args_list[1])))
  61.  
  62. if args_list[0] == "pulse_count" or args_list[0] == "count":
  63. self.faulty_worker.set_pulse_count(is_valid_number(int(args_list[1])))
  64.  
  65. if args_list[0] == "port":
  66. self.faulty_worker.set_serial_port(args_list[1])
  67. if not self.faulty_worker.validate_serial_connection():
  68. typer.secho("Invalid serial port.", fg=typer.colors.BRIGHT_RED)
  69. return
  70.  
  71. self.do_config(args)
  72.  
  73. def do_start(self, args):
  74. """Start the FaultyCat."""
  75. print("Starting the FaultyCat...")
  76. self.faulty_worker.start_faulty_attack()
  77.  
  78. def do_exit(self, line):
  79. """Exit the CLI."""
  80. return True
Buy Me A Coffee