mirror of
https://github.com/calebstewart/pwncat.git
synced 2024-11-27 19:04:15 +01:00
101 lines
3.2 KiB
Python
101 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
|
from prompt_toolkit.input.ansi_escape_sequences import (
|
|
REVERSE_ANSI_SEQUENCES,
|
|
ANSI_SEQUENCES,
|
|
)
|
|
from prompt_toolkit.keys import ALL_KEYS, Keys
|
|
from typing import Any, Dict, List, Union
|
|
import commentjson as json
|
|
import ipaddress
|
|
import re
|
|
import os
|
|
|
|
|
|
def key_type(value: str) -> bytes:
|
|
""" Converts a key name to a ansi keycode. The value can either be a single
|
|
printable character or a named key from prompt_toolkit Keys """
|
|
if len(value) == 1:
|
|
return value.encode("utf-8")
|
|
if value not in ALL_KEYS:
|
|
raise ValueError(f"invalid key: {value}")
|
|
key = [key for key in Keys if key.value == value][0]
|
|
return REVERSE_ANSI_SEQUENCES[key].encode("utf-8")
|
|
|
|
|
|
class KeyType:
|
|
def __init__(self, name: str):
|
|
if len(name) == 1:
|
|
self.value = name.encode("utf-8")
|
|
else:
|
|
if name not in ALL_KEYS:
|
|
raise ValueError(f"{name}: invalid key")
|
|
key = [key for key in Keys if key.value == name][0]
|
|
self.value = REVERSE_ANSI_SEQUENCES[key].encode("utf-8")
|
|
self.name = name
|
|
|
|
def __repr__(self):
|
|
return f"Key(name={repr(self.name)})"
|
|
|
|
def __bytes__(self):
|
|
return self.value
|
|
|
|
|
|
def local_file_type(value: str) -> str:
|
|
""" Ensure the local file exists """
|
|
if not os.path.isfile(value):
|
|
raise ValueError(f"{value}: no such file or directory")
|
|
return value
|
|
|
|
|
|
class Config:
|
|
def __init__(self):
|
|
|
|
# Basic key-value store w/ typing
|
|
self.values: Dict[str, Dict[str, Any]] = {
|
|
"lhost": {
|
|
"value": ipaddress.ip_address("127.0.0.1"),
|
|
"type": ipaddress.ip_address,
|
|
},
|
|
"prefix": {"value": KeyType("c-k"), "type": KeyType},
|
|
"privkey": {"value": "data/pwncat", "type": local_file_type},
|
|
"backdoor_user": {"value": "pwncat", "type": str},
|
|
"backdoor_pass": {"value": "pwncat", "type": str},
|
|
"on_load": {"value": "", "type": str},
|
|
"db": {"value": "sqlite:///:memory:", "type": str},
|
|
"cross": {"value": None, "type": str},
|
|
}
|
|
|
|
# Map ascii escape sequences or printable bytes to lists of commands to
|
|
# run.
|
|
self.bindings: Dict[KeyType, str] = {
|
|
KeyType("c-d"): "pass",
|
|
KeyType("s"): "sync",
|
|
KeyType("c"): "set state command",
|
|
}
|
|
|
|
def binding(self, name_or_value: Union[str, bytes]) -> str:
|
|
""" Get a key binding by it's key name or key value. """
|
|
|
|
if isinstance(name_or_value, bytes):
|
|
binding = [
|
|
b for key, b in self.bindings.items() if key.value == name_or_value
|
|
]
|
|
if not binding:
|
|
raise KeyError("no such key binding")
|
|
return binding[0]
|
|
|
|
key = KeyType(name_or_value)
|
|
return self.bindings[key]
|
|
|
|
def __getitem__(self, name: str) -> Any:
|
|
""" Get a configuration item """
|
|
return self.values[name]["value"]
|
|
|
|
def __setitem__(self, name: str, value: Any):
|
|
""" Set a configuration item """
|
|
item = self.values[name]
|
|
item["value"] = item["type"](value)
|
|
|
|
def __iter__(self):
|
|
yield from self.values
|