feat(auth,ai): add Personal Access Tokens support and persist
AI copilot context state
- Implement Personal Access Token (PAT) management in UserService with SHA-256
hash storage and O(1) lookup.
- Add gRPC endpoints (CreateApiToken, ListApiTokens, RevokeApiToken) and CLI
commands.
- Allow authentication using CONNPY_TOKEN environment variable in
ServiceProvider.
- Persist AI Copilot context mode and accumulation ranges across prompt
sessions in terminal_ui.
- Propagate node_info_json metadata over Copilot gRPC tunnel stream.
- Add unit tests for PAT lifecycle and Copilot context state persistence (430
passing tests).
This commit is contained in:
+3
-2
@@ -1159,8 +1159,9 @@ class ai:
|
||||
for msg in chat_history[-self.max_history:]:
|
||||
if msg.get('role') != 'system':
|
||||
messages.append(msg)
|
||||
# Add current user request
|
||||
messages.append({"role": "user", "content": clean_input})
|
||||
# Add current user request with a system note to prevent infinite escalation loops
|
||||
fallback_msg = clean_input + "\n\n[SYSTEM NOTE: The Architect is currently unavailable/failed to respond. You must handle the user's request directly as the Network Engineer. Do NOT attempt to escalate to the Architect again.]"
|
||||
messages.append({"role": "user", "content": fallback_msg})
|
||||
continue
|
||||
else:
|
||||
return {"response": f"Error: Both engines failed. {str(e)}", "chat_history": messages[1:], "usage": usage}
|
||||
|
||||
@@ -19,6 +19,14 @@ class LoginHandler:
|
||||
sys.exit(1)
|
||||
|
||||
def login(self, args):
|
||||
# Handle token management actions first
|
||||
if getattr(args, "create_token", None):
|
||||
return self.create_token(args)
|
||||
if getattr(args, "list_tokens", False):
|
||||
return self.list_tokens(args)
|
||||
if getattr(args, "revoke_token", None):
|
||||
return self.revoke_token(args)
|
||||
|
||||
if getattr(args, "status", False):
|
||||
return self.show_status()
|
||||
|
||||
@@ -141,3 +149,97 @@ class LoginHandler:
|
||||
printer.info(f"Expires at: {exp_dt.strftime('%Y-%m-%d %H:%M:%S UTC')}")
|
||||
except Exception as e:
|
||||
printer.error(f"Failed to check local session status: {e}")
|
||||
|
||||
def _get_auth_service(self):
|
||||
"""Gets an authenticated auth service stub, reusing existing or creating one."""
|
||||
auth_service = getattr(self.app.services, "auth", None)
|
||||
if not auth_service:
|
||||
import grpc
|
||||
from ..grpc_layer.stubs import AuthStub
|
||||
remote_host = self.app.services.remote_host or self.app.config.config.get("remote_host")
|
||||
if not remote_host:
|
||||
printer.error("Remote host is not configured. Run 'connpy config --remote HOST:PORT' first.")
|
||||
sys.exit(1)
|
||||
try:
|
||||
# Load existing session token for authentication
|
||||
token_path = os.path.join(self.app.config.defaultdir, ".token")
|
||||
if not os.path.exists(token_path):
|
||||
printer.error("No active session. Please log in first using 'connpy login'.")
|
||||
sys.exit(1)
|
||||
with open(token_path, "r") as f:
|
||||
session_token = f.read().strip()
|
||||
|
||||
from ..grpc_layer.stubs import AuthClientInterceptor
|
||||
interceptor = AuthClientInterceptor(lambda: session_token)
|
||||
channel = grpc.intercept_channel(grpc.insecure_channel(remote_host), interceptor)
|
||||
auth_service = AuthStub(channel, remote_host=remote_host)
|
||||
except Exception as e:
|
||||
printer.error(f"Failed to connect to remote server: {e}")
|
||||
sys.exit(1)
|
||||
return auth_service
|
||||
|
||||
def create_token(self, args):
|
||||
auth_service = self._get_auth_service()
|
||||
name = args.create_token
|
||||
expires_days = getattr(args, "expires_days", 0) or 0
|
||||
|
||||
try:
|
||||
result = auth_service.create_api_token(name, expires_in_days=expires_days)
|
||||
printer.success(f"API token '{name}' created successfully.")
|
||||
printer.warning("⚠ Copy this token now. It will NOT be shown again:")
|
||||
printer.data("Token", result["raw_token"])
|
||||
printer.info(f"Token ID: {result['token_id']}")
|
||||
if expires_days > 0:
|
||||
printer.info(f"Expires in: {expires_days} days")
|
||||
else:
|
||||
printer.info("Expires: Never (permanent)")
|
||||
except ConnpyError as e:
|
||||
printer.error(f"Failed to create token: {e}")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
printer.error(f"Failed to create token: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def list_tokens(self, args):
|
||||
auth_service = self._get_auth_service()
|
||||
|
||||
try:
|
||||
tokens = auth_service.list_api_tokens()
|
||||
if not tokens:
|
||||
printer.info("No API tokens found.")
|
||||
return
|
||||
|
||||
import yaml
|
||||
# Clean up empty strings from protobuf defaults
|
||||
cleaned = []
|
||||
for t in tokens:
|
||||
cleaned.append({
|
||||
"token_id": t["token_id"],
|
||||
"name": t["name"],
|
||||
"prefix": t["token_prefix"],
|
||||
"created": t["created_at"] or "N/A",
|
||||
"last_used": t["last_used_at"] or "Never",
|
||||
"expires": t["expires_at"] or "Never",
|
||||
})
|
||||
yaml_str = yaml.dump(cleaned, sort_keys=False, default_flow_style=False)
|
||||
printer.data("API Tokens", yaml_str)
|
||||
except ConnpyError as e:
|
||||
printer.error(f"Failed to list tokens: {e}")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
printer.error(f"Failed to list tokens: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def revoke_token(self, args):
|
||||
auth_service = self._get_auth_service()
|
||||
token_id = args.revoke_token
|
||||
|
||||
try:
|
||||
auth_service.revoke_api_token(token_id)
|
||||
printer.success(f"Token '{token_id}' revoked successfully.")
|
||||
except ConnpyError as e:
|
||||
printer.error(f"Failed to revoke token: {e}")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
printer.error(f"Failed to revoke token: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
+64
-13
@@ -30,20 +30,32 @@ class CopilotInterface:
|
||||
self.pt_input = pt_input
|
||||
self.pt_output = pt_output
|
||||
self.ai_service = AIService(config)
|
||||
self.session_state = session_state if session_state is not None else {
|
||||
'persona': 'engineer',
|
||||
'trust_mode': False,
|
||||
'memories': [],
|
||||
'os': None,
|
||||
'prompt': None
|
||||
}
|
||||
self.mode_range, self.mode_single, self.mode_lines = 0, 1, 2
|
||||
|
||||
self.session_state = session_state if session_state is not None else {}
|
||||
self.session_state.setdefault('persona', 'engineer')
|
||||
self.session_state.setdefault('trust_mode', False)
|
||||
self.session_state.setdefault('memories', [])
|
||||
self.session_state.setdefault('os', None)
|
||||
self.session_state.setdefault('prompt', None)
|
||||
self.session_state.setdefault('context_mode', self.mode_range)
|
||||
self.session_state.setdefault('context_cmd', 1)
|
||||
self.session_state.setdefault('context_lines', 50)
|
||||
self.session_state.setdefault('last_total_cmds', None)
|
||||
self.session_state.setdefault('last_total_lines', None)
|
||||
|
||||
if rich_file:
|
||||
self.console = Console(theme=connpy_theme, force_terminal=True, file=rich_file)
|
||||
else:
|
||||
self.console = Console(theme=connpy_theme)
|
||||
|
||||
self.mode_range, self.mode_single, self.mode_lines = 0, 1, 2
|
||||
def _sync_session_context(self, state: dict):
|
||||
"""Persist current context mode, depth, total commands, and total lines into session_state."""
|
||||
self.session_state['context_mode'] = state['context_mode']
|
||||
self.session_state['context_cmd'] = state['context_cmd']
|
||||
self.session_state['context_lines'] = state['context_lines']
|
||||
self.session_state['last_total_cmds'] = state['total_cmds']
|
||||
self.session_state['last_total_lines'] = state['total_lines']
|
||||
|
||||
def _get_theme_color(self, style_name: str, fallback: str = "white") -> str:
|
||||
"""Extract Hex or ANSI color name from the active rich theme."""
|
||||
@@ -77,16 +89,52 @@ class CopilotInterface:
|
||||
last_line = buffer.split('\n')[-1].strip() if buffer.strip() else "(prompt)"
|
||||
blocks = self.ai_service.build_context_blocks(raw_bytes, cmd_byte_positions, node_info, last_line=last_line)
|
||||
|
||||
total_cmds = len(blocks)
|
||||
total_lines = len(buffer.split('\n'))
|
||||
|
||||
saved_mode = self.session_state.get('context_mode', self.mode_range)
|
||||
saved_cmd = self.session_state.get('context_cmd', 1)
|
||||
saved_lines = self.session_state.get('context_lines', min(50, total_lines))
|
||||
last_total_cmds = self.session_state.get('last_total_cmds', None)
|
||||
last_total_lines = self.session_state.get('last_total_lines', None)
|
||||
|
||||
is_range = saved_mode in (self.mode_range, 0, 'RANGE', 'range')
|
||||
is_lines = saved_mode in (self.mode_lines, 2, 'LINES', 'lines')
|
||||
is_single = saved_mode in (self.mode_single, 1, 'SINGLE', 'single')
|
||||
|
||||
if is_range or is_single:
|
||||
if last_total_cmds is not None and total_cmds > last_total_cmds and saved_cmd > 1:
|
||||
new_cmds = total_cmds - last_total_cmds
|
||||
initial_cmd = saved_cmd + new_cmds
|
||||
else:
|
||||
initial_cmd = saved_cmd
|
||||
initial_lines = saved_lines
|
||||
elif is_lines:
|
||||
if last_total_lines is not None and total_lines > last_total_lines and saved_lines > 50:
|
||||
new_lines = total_lines - last_total_lines
|
||||
initial_lines = saved_lines + new_lines
|
||||
else:
|
||||
initial_lines = saved_lines
|
||||
initial_cmd = saved_cmd
|
||||
else:
|
||||
initial_cmd = saved_cmd
|
||||
initial_lines = saved_lines
|
||||
|
||||
state = {
|
||||
'context_cmd': 1,
|
||||
'total_cmds': len(blocks),
|
||||
'total_lines': len(buffer.split('\n')),
|
||||
'context_lines': min(50, len(buffer.split('\n'))),
|
||||
'context_mode': self.mode_range,
|
||||
'context_cmd': min(max(1, initial_cmd), max(1, total_cmds)),
|
||||
'total_cmds': total_cmds,
|
||||
'total_lines': total_lines,
|
||||
'context_lines': min(max(1, initial_lines), max(1, total_lines)),
|
||||
'context_mode': saved_mode,
|
||||
'cancelled': False,
|
||||
'toolbar_msg': '',
|
||||
'msg_expiry': 0
|
||||
}
|
||||
self.session_state['context_mode'] = saved_mode
|
||||
self.session_state['context_cmd'] = max(1, initial_cmd)
|
||||
self.session_state['context_lines'] = max(1, initial_lines)
|
||||
self.session_state['last_total_cmds'] = total_cmds
|
||||
self.session_state['last_total_lines'] = total_lines
|
||||
|
||||
# 1. Visual Separation
|
||||
self.console.print("") # Real line break
|
||||
@@ -105,6 +153,7 @@ class CopilotInterface:
|
||||
state['context_lines'] = min(state['context_lines'] + 50, state['total_lines'])
|
||||
else:
|
||||
state['context_cmd'] = min(state['context_cmd'] + 1, state['total_cmds'])
|
||||
self._sync_session_context(state)
|
||||
event.app.invalidate()
|
||||
@bindings.add('c-down')
|
||||
def _(event):
|
||||
@@ -112,6 +161,7 @@ class CopilotInterface:
|
||||
state['context_lines'] = max(state['context_lines'] - 50, min(50, state['total_lines']))
|
||||
else:
|
||||
state['context_cmd'] = max(state['context_cmd'] - 1, 1)
|
||||
self._sync_session_context(state)
|
||||
event.app.invalidate()
|
||||
@bindings.add('tab')
|
||||
def _(event):
|
||||
@@ -121,6 +171,7 @@ class CopilotInterface:
|
||||
buf.complete_next()
|
||||
else:
|
||||
state['context_mode'] = (state['context_mode'] + 1) % 3
|
||||
self._sync_session_context(state)
|
||||
event.app.invalidate()
|
||||
@bindings.add('escape', eager=True)
|
||||
@bindings.add('c-c')
|
||||
|
||||
@@ -368,7 +368,12 @@ def _build_tree(nodes, folders, profiles, plugins, configdir):
|
||||
},
|
||||
"user": user_dict,
|
||||
"sso": sso_dict,
|
||||
"login": {"--help": None, "-h": None, "*": None},
|
||||
"login": {
|
||||
"--status": None, "-s": None,
|
||||
"--create-token": None, "--list-tokens": None,
|
||||
"--revoke-token": None, "--expires-days": None,
|
||||
"--help": None, "-h": None, "*": None
|
||||
},
|
||||
"logout": {"--help": None, "-h": None},
|
||||
"config": config_dict,
|
||||
"sync": {
|
||||
|
||||
@@ -395,6 +395,10 @@ class connapp:
|
||||
loginparser.error = self._custom_error
|
||||
loginparser.add_argument("username", nargs='?', default=None, help="Username to authenticate")
|
||||
loginparser.add_argument("-s", "--status", action="store_true", help="Check current login status")
|
||||
loginparser.add_argument("--create-token", dest="create_token", metavar="NAME", help="Create a permanent API token with the given name")
|
||||
loginparser.add_argument("--list-tokens", dest="list_tokens", action="store_true", help="List all active API tokens")
|
||||
loginparser.add_argument("--revoke-token", dest="revoke_token", metavar="TOKEN_ID", help="Revoke an API token by its ID")
|
||||
loginparser.add_argument("--expires-days", dest="expires_days", type=int, default=0, metavar="DAYS", help="Optional expiration in days for --create-token (default: permanent)")
|
||||
loginparser.set_defaults(func=self._login.dispatch, action="login")
|
||||
|
||||
#LOGOUTPARSER
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -2652,6 +2652,21 @@ class AuthServiceStub(object):
|
||||
request_serializer=google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString,
|
||||
response_deserializer=connpy__pb2.SSOProvidersResponse.FromString,
|
||||
_registered_method=True)
|
||||
self.create_api_token = channel.unary_unary(
|
||||
'/connpy.AuthService/create_api_token',
|
||||
request_serializer=connpy__pb2.CreateApiTokenRequest.SerializeToString,
|
||||
response_deserializer=connpy__pb2.CreateApiTokenResponse.FromString,
|
||||
_registered_method=True)
|
||||
self.list_api_tokens = channel.unary_unary(
|
||||
'/connpy.AuthService/list_api_tokens',
|
||||
request_serializer=google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString,
|
||||
response_deserializer=connpy__pb2.ListApiTokensResponse.FromString,
|
||||
_registered_method=True)
|
||||
self.revoke_api_token = channel.unary_unary(
|
||||
'/connpy.AuthService/revoke_api_token',
|
||||
request_serializer=connpy__pb2.RevokeApiTokenRequest.SerializeToString,
|
||||
response_deserializer=google_dot_protobuf_dot_empty__pb2.Empty.FromString,
|
||||
_registered_method=True)
|
||||
|
||||
|
||||
class AuthServiceServicer(object):
|
||||
@@ -2681,6 +2696,24 @@ class AuthServiceServicer(object):
|
||||
context.set_details('Method not implemented!')
|
||||
raise NotImplementedError('Method not implemented!')
|
||||
|
||||
def create_api_token(self, request, context):
|
||||
"""Missing associated documentation comment in .proto file."""
|
||||
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
|
||||
context.set_details('Method not implemented!')
|
||||
raise NotImplementedError('Method not implemented!')
|
||||
|
||||
def list_api_tokens(self, request, context):
|
||||
"""Missing associated documentation comment in .proto file."""
|
||||
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
|
||||
context.set_details('Method not implemented!')
|
||||
raise NotImplementedError('Method not implemented!')
|
||||
|
||||
def revoke_api_token(self, request, context):
|
||||
"""Missing associated documentation comment in .proto file."""
|
||||
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
|
||||
context.set_details('Method not implemented!')
|
||||
raise NotImplementedError('Method not implemented!')
|
||||
|
||||
|
||||
def add_AuthServiceServicer_to_server(servicer, server):
|
||||
rpc_method_handlers = {
|
||||
@@ -2704,6 +2737,21 @@ def add_AuthServiceServicer_to_server(servicer, server):
|
||||
request_deserializer=google_dot_protobuf_dot_empty__pb2.Empty.FromString,
|
||||
response_serializer=connpy__pb2.SSOProvidersResponse.SerializeToString,
|
||||
),
|
||||
'create_api_token': grpc.unary_unary_rpc_method_handler(
|
||||
servicer.create_api_token,
|
||||
request_deserializer=connpy__pb2.CreateApiTokenRequest.FromString,
|
||||
response_serializer=connpy__pb2.CreateApiTokenResponse.SerializeToString,
|
||||
),
|
||||
'list_api_tokens': grpc.unary_unary_rpc_method_handler(
|
||||
servicer.list_api_tokens,
|
||||
request_deserializer=google_dot_protobuf_dot_empty__pb2.Empty.FromString,
|
||||
response_serializer=connpy__pb2.ListApiTokensResponse.SerializeToString,
|
||||
),
|
||||
'revoke_api_token': grpc.unary_unary_rpc_method_handler(
|
||||
servicer.revoke_api_token,
|
||||
request_deserializer=connpy__pb2.RevokeApiTokenRequest.FromString,
|
||||
response_serializer=google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString,
|
||||
),
|
||||
}
|
||||
generic_handler = grpc.method_handlers_generic_handler(
|
||||
'connpy.AuthService', rpc_method_handlers)
|
||||
@@ -2822,3 +2870,84 @@ class AuthService(object):
|
||||
timeout,
|
||||
metadata,
|
||||
_registered_method=True)
|
||||
|
||||
@staticmethod
|
||||
def create_api_token(request,
|
||||
target,
|
||||
options=(),
|
||||
channel_credentials=None,
|
||||
call_credentials=None,
|
||||
insecure=False,
|
||||
compression=None,
|
||||
wait_for_ready=None,
|
||||
timeout=None,
|
||||
metadata=None):
|
||||
return grpc.experimental.unary_unary(
|
||||
request,
|
||||
target,
|
||||
'/connpy.AuthService/create_api_token',
|
||||
connpy__pb2.CreateApiTokenRequest.SerializeToString,
|
||||
connpy__pb2.CreateApiTokenResponse.FromString,
|
||||
options,
|
||||
channel_credentials,
|
||||
insecure,
|
||||
call_credentials,
|
||||
compression,
|
||||
wait_for_ready,
|
||||
timeout,
|
||||
metadata,
|
||||
_registered_method=True)
|
||||
|
||||
@staticmethod
|
||||
def list_api_tokens(request,
|
||||
target,
|
||||
options=(),
|
||||
channel_credentials=None,
|
||||
call_credentials=None,
|
||||
insecure=False,
|
||||
compression=None,
|
||||
wait_for_ready=None,
|
||||
timeout=None,
|
||||
metadata=None):
|
||||
return grpc.experimental.unary_unary(
|
||||
request,
|
||||
target,
|
||||
'/connpy.AuthService/list_api_tokens',
|
||||
google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString,
|
||||
connpy__pb2.ListApiTokensResponse.FromString,
|
||||
options,
|
||||
channel_credentials,
|
||||
insecure,
|
||||
call_credentials,
|
||||
compression,
|
||||
wait_for_ready,
|
||||
timeout,
|
||||
metadata,
|
||||
_registered_method=True)
|
||||
|
||||
@staticmethod
|
||||
def revoke_api_token(request,
|
||||
target,
|
||||
options=(),
|
||||
channel_credentials=None,
|
||||
call_credentials=None,
|
||||
insecure=False,
|
||||
compression=None,
|
||||
wait_for_ready=None,
|
||||
timeout=None,
|
||||
metadata=None):
|
||||
return grpc.experimental.unary_unary(
|
||||
request,
|
||||
target,
|
||||
'/connpy.AuthService/revoke_api_token',
|
||||
connpy__pb2.RevokeApiTokenRequest.SerializeToString,
|
||||
google_dot_protobuf_dot_empty__pb2.Empty.FromString,
|
||||
options,
|
||||
channel_credentials,
|
||||
insecure,
|
||||
call_credentials,
|
||||
compression,
|
||||
wait_for_ready,
|
||||
timeout,
|
||||
metadata,
|
||||
_registered_method=True)
|
||||
|
||||
+124
-8
@@ -249,10 +249,59 @@ class NodeServicer(connpy_pb2_grpc.NodeServiceServicer):
|
||||
raw_bytes = str(raw_bytes).encode()
|
||||
|
||||
from connpy.utils import log_cleaner
|
||||
last_line = log_cleaner(raw_bytes.decode(errors='replace')).split('\n')[-1].strip()
|
||||
cleaned_buffer = log_cleaner(raw_bytes.decode(errors='replace'))
|
||||
last_line = cleaned_buffer.split('\n')[-1].strip() if cleaned_buffer.strip() else "(prompt)"
|
||||
blocks = service.build_context_blocks(raw_bytes, n.cmd_byte_positions, node_info, last_line=last_line)
|
||||
node_info["context_blocks"] = blocks
|
||||
|
||||
total_cmds = len(blocks)
|
||||
total_lines = len(cleaned_buffer.split('\n'))
|
||||
|
||||
if not hasattr(remote_stream, 'copilot_state') or remote_stream.copilot_state is None:
|
||||
remote_stream.copilot_state = {}
|
||||
session_state = remote_stream.copilot_state
|
||||
|
||||
if isinstance(node_info, dict):
|
||||
for k, v in node_info.items():
|
||||
if k in ('context_mode', 'context_cmd', 'context_lines', 'persona', 'trust', 'os', 'prompt'):
|
||||
session_state[k] = v
|
||||
|
||||
saved_mode = session_state.get('context_mode', 0)
|
||||
saved_cmd = session_state.get('context_cmd', 1)
|
||||
saved_lines = session_state.get('context_lines', 50)
|
||||
last_total_cmds = session_state.get('last_total_cmds', None)
|
||||
last_total_lines = session_state.get('last_total_lines', None)
|
||||
|
||||
is_range = saved_mode in (0, 'RANGE', 'range')
|
||||
is_lines = saved_mode in (2, 'LINES', 'lines')
|
||||
is_single = saved_mode in (1, 'SINGLE', 'single')
|
||||
|
||||
if is_range or is_single:
|
||||
if last_total_cmds is not None and total_cmds > last_total_cmds and saved_cmd > 1:
|
||||
new_cmds = total_cmds - last_total_cmds
|
||||
initial_cmd = saved_cmd + new_cmds
|
||||
else:
|
||||
initial_cmd = saved_cmd
|
||||
initial_lines = saved_lines
|
||||
elif is_lines:
|
||||
if last_total_lines is not None and total_lines > last_total_lines and saved_lines > 50:
|
||||
new_lines = total_lines - last_total_lines
|
||||
initial_lines = saved_lines + new_lines
|
||||
else:
|
||||
initial_lines = saved_lines
|
||||
initial_cmd = saved_cmd
|
||||
else:
|
||||
initial_cmd = saved_cmd
|
||||
initial_lines = saved_lines
|
||||
|
||||
session_state['context_cmd'] = max(1, initial_cmd)
|
||||
session_state['context_lines'] = max(1, initial_lines)
|
||||
session_state['last_total_cmds'] = total_cmds
|
||||
session_state['last_total_lines'] = total_lines
|
||||
|
||||
node_info.update(session_state)
|
||||
node_info['context_cmd'] = min(session_state['context_cmd'], max(1, total_cmds))
|
||||
node_info['context_lines'] = min(session_state['context_lines'], max(1, total_lines))
|
||||
node_info_json = json.dumps(node_info)
|
||||
|
||||
# Convert buffer to string if it's bytes for the preview
|
||||
@@ -297,6 +346,17 @@ class NodeServicer(connpy_pb2_grpc.NodeServiceServicer):
|
||||
if req_session_id and req_session_id != copilot_session_id:
|
||||
continue # Ignore stale request from a previous session
|
||||
|
||||
merged_node_info_str = req_data.get("node_info_json", "")
|
||||
if merged_node_info_str:
|
||||
try:
|
||||
merged_node_info = json.loads(merged_node_info_str)
|
||||
node_info.update(merged_node_info)
|
||||
# Sync context state from frontend into session_state for persistence
|
||||
for k in ('context_mode', 'context_cmd', 'context_lines'):
|
||||
if k in merged_node_info:
|
||||
session_state[k] = merged_node_info[k]
|
||||
except: pass
|
||||
|
||||
if "question" not in req_data or not req_data["question"] or req_data["question"] == "CANCEL" or req_data.get("action") in ("cancel", "web_cancel"):
|
||||
if req_data.get("action") == "web_cancel":
|
||||
os.write(child_fd, b'\x05')
|
||||
@@ -305,13 +365,6 @@ class NodeServicer(connpy_pb2_grpc.NodeServiceServicer):
|
||||
return
|
||||
question = req_data["question"]
|
||||
|
||||
merged_node_info_str = req_data.get("node_info_json", "")
|
||||
if merged_node_info_str:
|
||||
try:
|
||||
merged_node_info = json.loads(merged_node_info_str)
|
||||
node_info.update(merged_node_info)
|
||||
except: pass
|
||||
|
||||
context_buffer = req_data.get("context_buffer", "")
|
||||
if context_buffer.startswith('{"context_start_pos"'):
|
||||
try:
|
||||
@@ -373,6 +426,15 @@ class NodeServicer(connpy_pb2_grpc.NodeServiceServicer):
|
||||
if not action_data: return
|
||||
action = action_data.get("action", "cancel")
|
||||
|
||||
merged_node_info_str = action_data.get("node_info_json", "")
|
||||
if merged_node_info_str:
|
||||
try:
|
||||
merged_node_info = json.loads(merged_node_info_str)
|
||||
for k in ('context_mode', 'context_cmd', 'context_lines'):
|
||||
if k in merged_node_info:
|
||||
session_state[k] = merged_node_info[k]
|
||||
except: pass
|
||||
|
||||
if action == "continue":
|
||||
continue # Loop back for next question
|
||||
|
||||
@@ -1426,6 +1488,58 @@ class AuthServicer(connpy_pb2_grpc.AuthServiceServicer):
|
||||
|
||||
return Empty()
|
||||
|
||||
@handle_errors
|
||||
def create_api_token(self, request, context):
|
||||
username = _current_user.get()
|
||||
if not username:
|
||||
context.abort(grpc.StatusCode.UNAUTHENTICATED, "Authentication required")
|
||||
|
||||
try:
|
||||
expires_in_days = request.expires_in_days if request.expires_in_days > 0 else None
|
||||
result = self.registry.user_service.create_api_token(
|
||||
username, request.name, expires_in_days=expires_in_days
|
||||
)
|
||||
except ValueError as e:
|
||||
context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(e))
|
||||
|
||||
return connpy_pb2.CreateApiTokenResponse(
|
||||
token_id=result["token_id"],
|
||||
raw_token=result["raw_token"],
|
||||
name=result["name"],
|
||||
)
|
||||
|
||||
@handle_errors
|
||||
def list_api_tokens(self, request, context):
|
||||
username = _current_user.get()
|
||||
if not username:
|
||||
context.abort(grpc.StatusCode.UNAUTHENTICATED, "Authentication required")
|
||||
|
||||
tokens = self.registry.user_service.list_api_tokens(username)
|
||||
token_infos = [
|
||||
connpy_pb2.ApiTokenInfo(
|
||||
token_id=t["token_id"],
|
||||
name=t.get("name") or "",
|
||||
token_prefix=t.get("token_prefix") or "",
|
||||
created_at=t.get("created_at") or "",
|
||||
last_used_at=t.get("last_used_at") or "",
|
||||
expires_at=t.get("expires_at") or "",
|
||||
)
|
||||
for t in tokens
|
||||
]
|
||||
return connpy_pb2.ListApiTokensResponse(tokens=token_infos)
|
||||
|
||||
@handle_errors
|
||||
def revoke_api_token(self, request, context):
|
||||
username = _current_user.get()
|
||||
if not username:
|
||||
context.abort(grpc.StatusCode.UNAUTHENTICATED, "Authentication required")
|
||||
|
||||
removed = self.registry.user_service.revoke_api_token(username, request.token_id)
|
||||
if not removed:
|
||||
context.abort(grpc.StatusCode.NOT_FOUND, f"Token '{request.token_id}' not found")
|
||||
|
||||
return Empty()
|
||||
|
||||
class AuthInterceptor(grpc.ServerInterceptor):
|
||||
OPEN_METHODS = ["/connpy.AuthService/login", "/connpy.AuthService/login_sso", "/connpy.AuthService/get_sso_providers"]
|
||||
|
||||
@@ -1445,6 +1559,8 @@ class AuthInterceptor(grpc.ServerInterceptor):
|
||||
return self._unauthenticated_handler(handler_call_details, "Authorization token is missing")
|
||||
|
||||
username = self.registry.user_service.verify_jwt(token)
|
||||
if not username and token.startswith("cnp_pat_"):
|
||||
username = self.registry.user_service.verify_api_token(token)
|
||||
if not username:
|
||||
return self._unauthenticated_handler(handler_call_details, "Invalid or expired token")
|
||||
|
||||
|
||||
@@ -1148,3 +1148,33 @@ class AuthStub:
|
||||
def change_password(self, old_password, new_password):
|
||||
req = connpy_pb2.ChangePasswordRequest(old_password=old_password, new_password=new_password)
|
||||
self.stub.change_password(req)
|
||||
|
||||
@handle_errors
|
||||
def create_api_token(self, name, expires_in_days=0):
|
||||
req = connpy_pb2.CreateApiTokenRequest(name=name, expires_in_days=expires_in_days)
|
||||
resp = self.stub.create_api_token(req)
|
||||
return {
|
||||
"token_id": resp.token_id,
|
||||
"raw_token": resp.raw_token,
|
||||
"name": resp.name,
|
||||
}
|
||||
|
||||
@handle_errors
|
||||
def list_api_tokens(self):
|
||||
resp = self.stub.list_api_tokens(Empty())
|
||||
return [
|
||||
{
|
||||
"token_id": t.token_id,
|
||||
"name": t.name,
|
||||
"token_prefix": t.token_prefix,
|
||||
"created_at": t.created_at,
|
||||
"last_used_at": t.last_used_at,
|
||||
"expires_at": t.expires_at,
|
||||
}
|
||||
for t in resp.tokens
|
||||
]
|
||||
|
||||
@handle_errors
|
||||
def revoke_api_token(self, token_id):
|
||||
req = connpy_pb2.RevokeApiTokenRequest(token_id=token_id)
|
||||
self.stub.revoke_api_token(req)
|
||||
|
||||
@@ -304,6 +304,9 @@ service AuthService {
|
||||
rpc login_sso (LoginSSORequest) returns (LoginResponse) {}
|
||||
rpc change_password (ChangePasswordRequest) returns (google.protobuf.Empty) {}
|
||||
rpc get_sso_providers (google.protobuf.Empty) returns (SSOProvidersResponse) {}
|
||||
rpc create_api_token (CreateApiTokenRequest) returns (CreateApiTokenResponse) {}
|
||||
rpc list_api_tokens (google.protobuf.Empty) returns (ListApiTokensResponse) {}
|
||||
rpc revoke_api_token (RevokeApiTokenRequest) returns (google.protobuf.Empty) {}
|
||||
}
|
||||
|
||||
message SSOProvidersResponse {
|
||||
@@ -332,6 +335,34 @@ message ChangePasswordRequest {
|
||||
string new_password = 2;
|
||||
}
|
||||
|
||||
message CreateApiTokenRequest {
|
||||
string name = 1;
|
||||
int32 expires_in_days = 2;
|
||||
}
|
||||
|
||||
message CreateApiTokenResponse {
|
||||
string token_id = 1;
|
||||
string raw_token = 2;
|
||||
string name = 3;
|
||||
}
|
||||
|
||||
message ApiTokenInfo {
|
||||
string token_id = 1;
|
||||
string name = 2;
|
||||
string token_prefix = 3;
|
||||
string created_at = 4;
|
||||
string last_used_at = 5;
|
||||
string expires_at = 6;
|
||||
}
|
||||
|
||||
message ListApiTokensResponse {
|
||||
repeated ApiTokenInfo tokens = 1;
|
||||
}
|
||||
|
||||
message RevokeApiTokenRequest {
|
||||
string token_id = 1;
|
||||
}
|
||||
|
||||
message AnalyzeRequest {
|
||||
google.protobuf.Struct results = 1;
|
||||
string query = 2;
|
||||
|
||||
@@ -69,6 +69,9 @@ class ServiceProvider:
|
||||
)
|
||||
|
||||
def get_token():
|
||||
env_token = os.environ.get("CONNPY_TOKEN")
|
||||
if env_token:
|
||||
return env_token
|
||||
token_path = os.path.join(self.config.defaultdir, ".token")
|
||||
if os.path.exists(token_path):
|
||||
try:
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import os
|
||||
import hashlib
|
||||
import re
|
||||
import shutil
|
||||
import secrets
|
||||
@@ -18,6 +19,9 @@ class UserService:
|
||||
# Ensure users directory exists
|
||||
os.makedirs(self.users_dir, exist_ok=True)
|
||||
|
||||
# Reverse index cache: token_hash -> (username, token_id)
|
||||
self._token_index: dict[str, tuple[str, str]] = {}
|
||||
|
||||
def _load_registry(self) -> dict:
|
||||
"""Loads registry from file. If it doesn't exist, initializes it with a new JWT secret."""
|
||||
if not os.path.exists(self.registry_file):
|
||||
@@ -61,6 +65,16 @@ class UserService:
|
||||
pass
|
||||
raise e
|
||||
|
||||
def _build_token_index(self, registry: dict) -> dict[str, tuple[str, str]]:
|
||||
"""Builds a reverse index of token_hash -> (username, token_id) for O(1) PAT lookup."""
|
||||
index = {}
|
||||
for username, user_data in registry.get("users", {}).items():
|
||||
for token_id, token_meta in user_data.get("api_tokens", {}).items():
|
||||
token_hash = token_meta.get("token_hash")
|
||||
if token_hash:
|
||||
index[token_hash] = (username, token_id)
|
||||
return index
|
||||
|
||||
def create_user(self, username, password, config_path=None) -> dict:
|
||||
"""Creates a new user with bcrypt-hashed credentials.
|
||||
|
||||
@@ -237,3 +251,125 @@ class UserService:
|
||||
return payload.get("sub")
|
||||
except (jwt.ExpiredSignatureError, jwt.InvalidTokenError, KeyError):
|
||||
return None
|
||||
|
||||
# --- Personal Access Token (PAT) Management ---
|
||||
|
||||
def create_api_token(self, username: str, name: str, expires_in_days: int | None = None) -> dict:
|
||||
"""Creates a Personal Access Token for the user.
|
||||
|
||||
Returns the raw token ONCE. Only the SHA-256 hash is persisted.
|
||||
"""
|
||||
if not name or not isinstance(name, str):
|
||||
raise ValueError("Token name cannot be empty")
|
||||
|
||||
registry = self._load_registry()
|
||||
if username not in registry["users"]:
|
||||
raise ValueError(f"User '{username}' not found")
|
||||
|
||||
user_data = registry["users"][username]
|
||||
if "api_tokens" not in user_data:
|
||||
user_data["api_tokens"] = {}
|
||||
|
||||
# Generate cryptographically secure token with recognizable prefix
|
||||
raw_secret = secrets.token_hex(32)
|
||||
raw_token = f"cnp_pat_{raw_secret}"
|
||||
token_hash = hashlib.sha256(raw_token.encode("utf-8")).hexdigest()
|
||||
token_id = f"tok_{secrets.token_hex(4)}"
|
||||
|
||||
now = datetime.datetime.now(datetime.timezone.utc)
|
||||
expires_at = None
|
||||
if expires_in_days and expires_in_days > 0:
|
||||
expires_at = (now + datetime.timedelta(days=expires_in_days)).isoformat()
|
||||
|
||||
user_data["api_tokens"][token_id] = {
|
||||
"name": name,
|
||||
"token_hash": token_hash,
|
||||
"token_prefix": raw_token[:16],
|
||||
"created_at": now.isoformat(),
|
||||
"last_used_at": None,
|
||||
"expires_at": expires_at,
|
||||
}
|
||||
|
||||
self._save_registry(registry)
|
||||
self._token_index = self._build_token_index(registry)
|
||||
|
||||
return {
|
||||
"token_id": token_id,
|
||||
"raw_token": raw_token,
|
||||
"name": name,
|
||||
}
|
||||
|
||||
def list_api_tokens(self, username: str) -> list[dict]:
|
||||
"""Lists all active API tokens for a user (without sensitive data)."""
|
||||
registry = self._load_registry()
|
||||
if username not in registry["users"]:
|
||||
raise ValueError(f"User '{username}' not found")
|
||||
|
||||
tokens = registry["users"][username].get("api_tokens", {})
|
||||
return [
|
||||
{
|
||||
"token_id": tid,
|
||||
"name": meta.get("name"),
|
||||
"token_prefix": meta.get("token_prefix"),
|
||||
"created_at": meta.get("created_at"),
|
||||
"last_used_at": meta.get("last_used_at"),
|
||||
"expires_at": meta.get("expires_at"),
|
||||
}
|
||||
for tid, meta in tokens.items()
|
||||
]
|
||||
|
||||
def revoke_api_token(self, username: str, token_id: str) -> bool:
|
||||
"""Revokes (deletes) a specific API token. Returns True if found and removed."""
|
||||
registry = self._load_registry()
|
||||
if username not in registry["users"]:
|
||||
raise ValueError(f"User '{username}' not found")
|
||||
|
||||
tokens = registry["users"][username].get("api_tokens", {})
|
||||
if token_id not in tokens:
|
||||
return False
|
||||
|
||||
del tokens[token_id]
|
||||
self._save_registry(registry)
|
||||
self._token_index = self._build_token_index(registry)
|
||||
return True
|
||||
|
||||
def verify_api_token(self, raw_token: str) -> str | None:
|
||||
"""Validates a PAT by hashing it and looking up the reverse index.
|
||||
|
||||
Returns username if valid and not expired, None otherwise.
|
||||
"""
|
||||
token_hash = hashlib.sha256(raw_token.encode("utf-8")).hexdigest()
|
||||
|
||||
# Rebuild index if empty (cold start or after process restart)
|
||||
if not self._token_index:
|
||||
registry = self._load_registry()
|
||||
self._token_index = self._build_token_index(registry)
|
||||
|
||||
match = self._token_index.get(token_hash)
|
||||
if not match:
|
||||
return None
|
||||
|
||||
username, token_id = match
|
||||
|
||||
# Validate token still exists and check expiration
|
||||
registry = self._load_registry()
|
||||
user_data = registry.get("users", {}).get(username, {})
|
||||
token_meta = user_data.get("api_tokens", {}).get(token_id)
|
||||
|
||||
if not token_meta:
|
||||
# Token was revoked between index build and now
|
||||
self._token_index = self._build_token_index(registry)
|
||||
return None
|
||||
|
||||
# Check expiration
|
||||
expires_at = token_meta.get("expires_at")
|
||||
if expires_at:
|
||||
exp_dt = datetime.datetime.fromisoformat(expires_at)
|
||||
if datetime.datetime.now(datetime.timezone.utc) > exp_dt:
|
||||
return None
|
||||
|
||||
# Update last_used_at
|
||||
token_meta["last_used_at"] = datetime.datetime.now(datetime.timezone.utc).isoformat()
|
||||
self._save_registry(registry)
|
||||
|
||||
return username
|
||||
|
||||
@@ -400,3 +400,164 @@ def test_build_context_blocks_pager_scrolling_6wind_escapes():
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
def test_copilot_context_state_persistence():
|
||||
from connpy.cli.terminal_ui import CopilotInterface
|
||||
|
||||
class MockConfig:
|
||||
def __init__(self):
|
||||
self.config = {"ai": {}}
|
||||
self.defaultdir = "/tmp"
|
||||
|
||||
session_state = {}
|
||||
interface = CopilotInterface(MockConfig(), session_state=session_state)
|
||||
|
||||
raw_bytes = b"router# show ip\r\nrouter# show run\r\nrouter# "
|
||||
blocks = [
|
||||
(0, 15, "router# show ip"),
|
||||
(15, 30, "router# show run"),
|
||||
(30, 40, "router#")
|
||||
]
|
||||
|
||||
async def mock_ai_call(active_buffer, question, on_chunk, node_info):
|
||||
return {"guide": "Ok", "commands": [], "risk_level": "low"}
|
||||
|
||||
async def mock_prompt_async(self, *args, **kwargs):
|
||||
kb = kwargs.get('key_bindings')
|
||||
if kb:
|
||||
class DummyApp:
|
||||
def invalidate(self): pass
|
||||
class DummyEvent:
|
||||
app = DummyApp()
|
||||
current_buffer = type('Buf', (), {'text': ''})()
|
||||
|
||||
# Trigger TAB key ('c-i' or 'tab') to switch mode from RANGE (0) to SINGLE (1)
|
||||
for b in kb.bindings:
|
||||
if any(k in ('c-i', 'tab') or 'tab' in str(k).lower() or 'c-i' in str(k).lower() for k in b.keys):
|
||||
b.handler(DummyEvent())
|
||||
break
|
||||
return "test question"
|
||||
|
||||
with patch('prompt_toolkit.PromptSession.prompt_async', mock_prompt_async):
|
||||
asyncio.run(interface.run_session(
|
||||
raw_bytes=raw_bytes,
|
||||
node_info={"name": "test"},
|
||||
on_ai_call=mock_ai_call,
|
||||
blocks=blocks
|
||||
))
|
||||
|
||||
assert interface.session_state.get('context_mode') == interface.mode_single
|
||||
assert interface.session_state.get('last_total_cmds') == len(blocks)
|
||||
|
||||
|
||||
def test_copilot_range_mode_accumulation():
|
||||
from connpy.cli.terminal_ui import CopilotInterface
|
||||
|
||||
class MockConfig:
|
||||
def __init__(self):
|
||||
self.config = {"ai": {}}
|
||||
self.defaultdir = "/tmp"
|
||||
|
||||
raw_bytes = b"router# cmd1\r\nrouter# cmd2\r\nrouter# cmd3\r\nrouter# "
|
||||
blocks = [
|
||||
(0, 10, "router# cmd1"),
|
||||
(10, 20, "router# cmd2"),
|
||||
(20, 30, "router# cmd3"),
|
||||
(30, 40, "router#")
|
||||
]
|
||||
|
||||
async def mock_ai_call(active_buffer, question, on_chunk, node_info):
|
||||
return {"guide": "Ok", "commands": [], "risk_level": "low"}
|
||||
|
||||
async def mock_prompt_async(self, *args, **kwargs):
|
||||
return "cancel"
|
||||
|
||||
# Test 1: RANGE mode at default (saved_cmd = 1) -> stays 1 (does not expand)
|
||||
session_state_default = {'context_mode': 0, 'context_cmd': 1, 'last_total_cmds': 2}
|
||||
interface_default = CopilotInterface(MockConfig(), session_state=session_state_default)
|
||||
with patch('prompt_toolkit.PromptSession.prompt_async', mock_prompt_async):
|
||||
asyncio.run(interface_default.run_session(raw_bytes=raw_bytes, node_info={"name": "test"}, on_ai_call=mock_ai_call, blocks=blocks))
|
||||
assert interface_default.session_state.get('context_cmd') == 1
|
||||
|
||||
# Test 2: RANGE mode at expanded (saved_cmd = 2 > 1) -> expands to 2 + 2 = 4
|
||||
session_state_expanded = {'context_mode': 0, 'context_cmd': 2, 'last_total_cmds': 2}
|
||||
interface_expanded = CopilotInterface(MockConfig(), session_state=session_state_expanded)
|
||||
with patch('prompt_toolkit.PromptSession.prompt_async', mock_prompt_async):
|
||||
asyncio.run(interface_expanded.run_session(raw_bytes=raw_bytes, node_info={"name": "test"}, on_ai_call=mock_ai_call, blocks=blocks))
|
||||
assert interface_expanded.session_state.get('context_cmd') == 4
|
||||
|
||||
|
||||
def test_copilot_lines_mode_accumulation():
|
||||
from connpy.cli.terminal_ui import CopilotInterface
|
||||
|
||||
class MockConfig:
|
||||
def __init__(self):
|
||||
self.config = {"ai": {}}
|
||||
self.defaultdir = "/tmp"
|
||||
|
||||
raw_bytes = ("line\n" * 130).encode()
|
||||
blocks = [(0, 10, "router#")]
|
||||
|
||||
async def mock_ai_call(active_buffer, question, on_chunk, node_info):
|
||||
return {"guide": "Ok", "commands": [], "risk_level": "low"}
|
||||
|
||||
async def mock_prompt_async(self, *args, **kwargs):
|
||||
return "cancel"
|
||||
|
||||
# Test 1: LINES mode at default 50 lines -> stays 50
|
||||
session_state_default = {'context_mode': 2, 'context_lines': 50, 'last_total_lines': 100}
|
||||
interface_default = CopilotInterface(MockConfig(), session_state=session_state_default)
|
||||
with patch('prompt_toolkit.PromptSession.prompt_async', mock_prompt_async):
|
||||
asyncio.run(interface_default.run_session(raw_bytes=raw_bytes, node_info={"name": "test"}, on_ai_call=mock_ai_call, blocks=blocks))
|
||||
assert interface_default.session_state.get('context_lines') == 50
|
||||
|
||||
# Test 2: LINES mode at expanded 100 lines -> expands beyond 100
|
||||
session_state_expanded = {'context_mode': 2, 'context_lines': 100, 'last_total_lines': 100}
|
||||
interface_expanded = CopilotInterface(MockConfig(), session_state=session_state_expanded)
|
||||
with patch('prompt_toolkit.PromptSession.prompt_async', mock_prompt_async):
|
||||
asyncio.run(interface_expanded.run_session(raw_bytes=raw_bytes, node_info={"name": "test"}, on_ai_call=mock_ai_call, blocks=blocks))
|
||||
assert interface_expanded.session_state.get('context_lines') > 100
|
||||
|
||||
|
||||
def test_copilot_single_mode_retains_command_block():
|
||||
from connpy.cli.terminal_ui import CopilotInterface
|
||||
|
||||
class MockConfig:
|
||||
def __init__(self):
|
||||
self.config = {"ai": {}}
|
||||
self.defaultdir = "/tmp"
|
||||
|
||||
raw_bytes = b"router# cmd1\r\nrouter# cmd2\r\nrouter# cmd3\r\nrouter# "
|
||||
blocks = [
|
||||
(0, 10, "router# cmd1"),
|
||||
(10, 20, "router# cmd2"),
|
||||
(20, 30, "router# cmd3"),
|
||||
(30, 40, "router#")
|
||||
]
|
||||
|
||||
async def mock_ai_call(active_buffer, question, on_chunk, node_info):
|
||||
return {"guide": "Ok", "commands": [], "risk_level": "low"}
|
||||
|
||||
async def mock_prompt_async(self, *args, **kwargs):
|
||||
return "cancel"
|
||||
|
||||
# Test 1: In SINGLE mode at default (context_cmd = 1), stays at 1
|
||||
session_state_default = {'context_mode': 1, 'context_cmd': 1, 'last_total_cmds': 2}
|
||||
interface_default = CopilotInterface(MockConfig(), session_state=session_state_default)
|
||||
with patch('prompt_toolkit.PromptSession.prompt_async', mock_prompt_async):
|
||||
asyncio.run(interface_default.run_session(raw_bytes=raw_bytes, node_info={"name": "test"}, on_ai_call=mock_ai_call, blocks=blocks))
|
||||
assert interface_default.session_state.get('context_cmd') == 1
|
||||
|
||||
# Test 2: In SINGLE mode at past command (context_cmd = 2 > 1), becomes 2 + 2 = 4 to stay locked on past command
|
||||
session_state_custom = {'context_mode': 1, 'context_cmd': 2, 'last_total_cmds': 2}
|
||||
interface_custom = CopilotInterface(MockConfig(), session_state=session_state_custom)
|
||||
with patch('prompt_toolkit.PromptSession.prompt_async', mock_prompt_async):
|
||||
asyncio.run(interface_custom.run_session(raw_bytes=raw_bytes, node_info={"name": "test"}, on_ai_call=mock_ai_call, blocks=blocks))
|
||||
assert interface_custom.session_state.get('context_cmd') == 4
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,216 @@
|
||||
import os
|
||||
import datetime
|
||||
import hashlib
|
||||
import pytest
|
||||
import yaml
|
||||
from connpy.services.user_service import UserService
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_config_dir(tmp_path):
|
||||
"""Creates a temporary config directory for testing."""
|
||||
config_dir = tmp_path / "conn_config"
|
||||
config_dir.mkdir()
|
||||
return config_dir
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def user_service(test_config_dir):
|
||||
"""Initializes UserService pointing to a temporary directory."""
|
||||
return UserService(str(test_config_dir))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def user_with_token(user_service):
|
||||
"""Creates a user and returns (user_service, username, token_result)."""
|
||||
username = "tokenuser"
|
||||
user_service.create_user(username, "password123")
|
||||
result = user_service.create_api_token(username, "Test Token")
|
||||
return user_service, username, result
|
||||
|
||||
|
||||
class TestApiTokenCreation:
|
||||
def test_create_api_token_returns_raw_token(self, user_service):
|
||||
"""Verifies that create_api_token returns a raw token with the correct prefix."""
|
||||
user_service.create_user("alice", "pass")
|
||||
result = user_service.create_api_token("alice", "CI Pipeline")
|
||||
|
||||
assert "raw_token" in result
|
||||
assert result["raw_token"].startswith("cnp_pat_")
|
||||
assert len(result["raw_token"]) > 16
|
||||
assert "token_id" in result
|
||||
assert result["token_id"].startswith("tok_")
|
||||
assert result["name"] == "CI Pipeline"
|
||||
|
||||
def test_create_api_token_stores_hash_not_plaintext(self, user_service):
|
||||
"""Ensures only the SHA-256 hash is persisted, never the raw token."""
|
||||
user_service.create_user("bob", "pass")
|
||||
result = user_service.create_api_token("bob", "My App")
|
||||
|
||||
registry = user_service._load_registry()
|
||||
tokens = registry["users"]["bob"]["api_tokens"]
|
||||
assert len(tokens) == 1
|
||||
|
||||
token_meta = list(tokens.values())[0]
|
||||
expected_hash = hashlib.sha256(result["raw_token"].encode("utf-8")).hexdigest()
|
||||
assert token_meta["token_hash"] == expected_hash
|
||||
# Raw token must NOT be stored
|
||||
assert result["raw_token"] not in str(token_meta)
|
||||
|
||||
def test_create_api_token_with_expiration(self, user_service):
|
||||
"""Verifies that expires_at is set correctly when expires_in_days is provided."""
|
||||
user_service.create_user("charlie", "pass")
|
||||
user_service.create_api_token("charlie", "Temp Token", expires_in_days=30)
|
||||
|
||||
registry = user_service._load_registry()
|
||||
token_meta = list(registry["users"]["charlie"]["api_tokens"].values())[0]
|
||||
assert token_meta["expires_at"] is not None
|
||||
|
||||
exp_dt = datetime.datetime.fromisoformat(token_meta["expires_at"])
|
||||
now = datetime.datetime.now(datetime.timezone.utc)
|
||||
delta = exp_dt - now
|
||||
assert 29 <= delta.days <= 30
|
||||
|
||||
def test_create_api_token_permanent_by_default(self, user_service):
|
||||
"""Verifies that expires_at is None when no expiration is specified."""
|
||||
user_service.create_user("dave", "pass")
|
||||
user_service.create_api_token("dave", "Permanent Token")
|
||||
|
||||
registry = user_service._load_registry()
|
||||
token_meta = list(registry["users"]["dave"]["api_tokens"].values())[0]
|
||||
assert token_meta["expires_at"] is None
|
||||
|
||||
def test_create_api_token_nonexistent_user(self, user_service):
|
||||
"""Ensures creating a token for a non-existent user raises ValueError."""
|
||||
with pytest.raises(ValueError, match="not found"):
|
||||
user_service.create_api_token("ghost", "Token")
|
||||
|
||||
def test_create_api_token_empty_name(self, user_service):
|
||||
"""Ensures empty token names are rejected."""
|
||||
user_service.create_user("eve", "pass")
|
||||
with pytest.raises(ValueError, match="cannot be empty"):
|
||||
user_service.create_api_token("eve", "")
|
||||
|
||||
def test_create_multiple_tokens(self, user_service):
|
||||
"""Verifies a user can have multiple tokens."""
|
||||
user_service.create_user("frank", "pass")
|
||||
t1 = user_service.create_api_token("frank", "Token 1")
|
||||
t2 = user_service.create_api_token("frank", "Token 2")
|
||||
|
||||
assert t1["token_id"] != t2["token_id"]
|
||||
assert t1["raw_token"] != t2["raw_token"]
|
||||
|
||||
tokens = user_service.list_api_tokens("frank")
|
||||
assert len(tokens) == 2
|
||||
|
||||
|
||||
class TestApiTokenVerification:
|
||||
def test_verify_valid_token(self, user_with_token):
|
||||
"""Verifies that a valid raw token authenticates correctly."""
|
||||
svc, username, result = user_with_token
|
||||
verified = svc.verify_api_token(result["raw_token"])
|
||||
assert verified == username
|
||||
|
||||
def test_verify_invalid_token(self, user_service):
|
||||
"""Verifies that a random/invalid token returns None."""
|
||||
user_service.create_user("alice", "pass")
|
||||
assert user_service.verify_api_token("cnp_pat_invalid_token_here") is None
|
||||
|
||||
def test_verify_expired_token(self, user_service):
|
||||
"""Verifies that an expired token returns None."""
|
||||
user_service.create_user("alice", "pass")
|
||||
result = user_service.create_api_token("alice", "Expiring", expires_in_days=1)
|
||||
|
||||
# Manually set expires_at to the past
|
||||
registry = user_service._load_registry()
|
||||
token_meta = list(registry["users"]["alice"]["api_tokens"].values())[0]
|
||||
token_meta["expires_at"] = (
|
||||
datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(hours=1)
|
||||
).isoformat()
|
||||
user_service._save_registry(registry)
|
||||
# Invalidate cache so verify_api_token re-reads
|
||||
user_service._token_index = {}
|
||||
|
||||
assert user_service.verify_api_token(result["raw_token"]) is None
|
||||
|
||||
def test_verify_updates_last_used_at(self, user_with_token):
|
||||
"""Verifies that last_used_at is updated upon successful verification."""
|
||||
svc, username, result = user_with_token
|
||||
|
||||
# Initially last_used_at should be None
|
||||
registry = svc._load_registry()
|
||||
token_meta = list(registry["users"][username]["api_tokens"].values())[0]
|
||||
assert token_meta["last_used_at"] is None
|
||||
|
||||
# Verify the token
|
||||
svc.verify_api_token(result["raw_token"])
|
||||
|
||||
# Now last_used_at should be set
|
||||
registry = svc._load_registry()
|
||||
token_meta = list(registry["users"][username]["api_tokens"].values())[0]
|
||||
assert token_meta["last_used_at"] is not None
|
||||
|
||||
|
||||
class TestApiTokenListing:
|
||||
def test_list_tokens_returns_metadata(self, user_with_token):
|
||||
"""Verifies list returns metadata without sensitive data."""
|
||||
svc, username, result = user_with_token
|
||||
tokens = svc.list_api_tokens(username)
|
||||
|
||||
assert len(tokens) == 1
|
||||
t = tokens[0]
|
||||
assert t["token_id"] == result["token_id"]
|
||||
assert t["name"] == "Test Token"
|
||||
assert t["token_prefix"].startswith("cnp_pat_")
|
||||
assert "created_at" in t
|
||||
# Must NOT expose token_hash or raw_token
|
||||
assert "token_hash" not in t
|
||||
assert "raw_token" not in t
|
||||
|
||||
def test_list_tokens_empty(self, user_service):
|
||||
"""Verifies listing tokens for a user with none returns empty list."""
|
||||
user_service.create_user("alice", "pass")
|
||||
assert user_service.list_api_tokens("alice") == []
|
||||
|
||||
def test_list_tokens_nonexistent_user(self, user_service):
|
||||
"""Ensures listing tokens for a non-existent user raises ValueError."""
|
||||
with pytest.raises(ValueError, match="not found"):
|
||||
user_service.list_api_tokens("ghost")
|
||||
|
||||
|
||||
class TestApiTokenRevocation:
|
||||
def test_revoke_token(self, user_with_token):
|
||||
"""Verifies that a revoked token is immediately invalid."""
|
||||
svc, username, result = user_with_token
|
||||
|
||||
# Token works before revocation
|
||||
assert svc.verify_api_token(result["raw_token"]) == username
|
||||
|
||||
# Revoke
|
||||
removed = svc.revoke_api_token(username, result["token_id"])
|
||||
assert removed is True
|
||||
|
||||
# Token must fail after revocation
|
||||
assert svc.verify_api_token(result["raw_token"]) is None
|
||||
|
||||
# List should be empty
|
||||
assert svc.list_api_tokens(username) == []
|
||||
|
||||
def test_revoke_nonexistent_token(self, user_service):
|
||||
"""Verifies revoking a non-existent token returns False."""
|
||||
user_service.create_user("alice", "pass")
|
||||
assert user_service.revoke_api_token("alice", "tok_nonexistent") is False
|
||||
|
||||
def test_revoke_nonexistent_user(self, user_service):
|
||||
"""Ensures revoking a token for a non-existent user raises ValueError."""
|
||||
with pytest.raises(ValueError, match="not found"):
|
||||
user_service.revoke_api_token("ghost", "tok_abc")
|
||||
|
||||
|
||||
class TestJwtUnchanged:
|
||||
def test_jwt_still_works(self, user_service):
|
||||
"""Confirms that existing JWT session tokens still authenticate correctly."""
|
||||
user_service.create_user("jwtuser", "pass")
|
||||
token = user_service.generate_jwt("jwtuser")
|
||||
verified = user_service.verify_jwt(token)
|
||||
assert verified == "jwtuser"
|
||||
@@ -172,6 +172,8 @@ class RemoteStream:
|
||||
})
|
||||
if getattr(req, "copilot_action", ""):
|
||||
copilot_msg["action"] = req.copilot_action
|
||||
if getattr(req, "copilot_node_info_json", ""):
|
||||
copilot_msg["node_info_json"] = req.copilot_node_info_json
|
||||
|
||||
if copilot_msg:
|
||||
self._loop.call_soon_threadsafe(self.copilot_queue.put_nowait, copilot_msg)
|
||||
|
||||
Reference in New Issue
Block a user