codecraftai-protocol-ssh 1.3.0__py3-none-any.whl

ccai/protocols/ssh/sshkeys.py

@@ -0,0 +1,197 @@
+"""
+SSH Keys Interface
+"""
+from abc import ABC, abstractmethod
+from cryptography.hazmat.primitives import serialization, hashes
+from cryptography.hazmat.primitives.asymmetric import rsa, ed25519, padding
+from cryptography.exceptions import InvalidSignature
+
+class HostKey(ABC):
+    """
+    Abstract base class for Host Keys.
+    """
+    
+    @abstractmethod
+    def get_name(self) -> str:
+        """
+        Return algorithm name (e.g. ssh-rsa)
+        """
+        pass
+        
+    @abstractmethod
+    def sign(self, data: bytes) -> bytes:
+        """
+        Sign data.
+        """
+        pass
+        
+    @abstractmethod
+    def verify(self, data: bytes, signature: bytes) -> bool:
+        """
+        Verify signature.
+        """
+        pass
+        
+    @abstractmethod
+    def as_bytes(self) -> bytes:
+        """
+        Return public key as bytes for transmission (OpenSSH format).
+        """
+        pass
+
+class SshRSAKey(HostKey):
+    """
+    RSA Key Implementation using cryptography.
+    """
+    def __init__(self, key: rsa.RSAPrivateKey = None, pub_key: rsa.RSAPublicKey = None):
+        if key:
+            self.private_key = key
+            self.public_key = key.public_key()
+        elif pub_key:
+            self.private_key = None
+            self.public_key = pub_key
+        else:
+            # Generate a new key for testing if none provided? 
+            # Or raise error? For now, let's generate 2048-bit key if nothing passed.
+            self.private_key = rsa.generate_private_key(
+                public_exponent=65537,
+                key_size=2048
+            )
+            self.public_key = self.private_key.public_key()
+            
+    @classmethod
+    def from_private_key_file(cls, filename, password=None):
+        with open(filename, "rb") as f:
+            data = f.read()
+        return cls.from_private_key(data, password)
+
+    @classmethod
+    def from_private_key(cls, data, password=None):
+        # Handle file-like object
+        if hasattr(data, "read"):
+            data = data.read()
+            
+        if isinstance(data, str):
+            data = data.encode('utf-8')
+            
+        p_key = serialization.load_pem_private_key(
+            data,
+            password=password.encode() if isinstance(password, str) else password
+        )
+        if not isinstance(p_key, rsa.RSAPrivateKey):
+            raise ValueError("Not an RSA key")
+            
+        return cls(key=p_key)
+        
+    def get_name(self) -> str:
+        return "ssh-rsa"
+        
+    def sign(self, data: bytes) -> bytes:
+        if not self.private_key:
+            raise ValueError("Private key not available for signing")
+            
+        return self.private_key.sign(
+            data,
+            padding.PKCS1v15(),
+            hashes.SHA1()
+        )
+
+    def verify(self, data: bytes, signature: bytes) -> bool:
+        try:
+            self.public_key.verify(
+                signature,
+                data,
+                padding.PKCS1v15(),
+                hashes.SHA1()
+            )
+            return True
+        except InvalidSignature:
+            return False
+
+    def as_bytes(self) -> bytes:
+        """
+        Format: string "ssh-rsa", mpint e, mpint n
+        """
+        # We need to manually construct the blob because cryptography's PublicFormat.OpenSSH 
+        # includes the header "ssh-rsa <b64> <comment>", but we just want the inner blob 
+        # or we want the wire format components.
+        # The wire format is: string algorithm, mpint e, mpint n.
+        
+        numbers = self.public_key.public_numbers()
+        
+        from ccai.protocols.ssh.sshpacket import SshPacket
+        pkt = SshPacket()
+        pkt.add_string(self.get_name())
+        pkt.add_mpint(numbers.e)
+        pkt.add_mpint(numbers.n)
+        return pkt.payload
+
+class SshEd25519Key(HostKey):
+    """
+    Ed25519 Key Implementation.
+    """
+    def __init__(self, key: ed25519.Ed25519PrivateKey = None, pub_key: ed25519.Ed25519PublicKey = None):
+        if key:
+            self.private_key = key
+            self.public_key = key.public_key()
+        elif pub_key:
+            self.private_key = None
+            self.public_key = pub_key
+        else:
+            self.private_key = ed25519.Ed25519PrivateKey.generate()
+            self.public_key = self.private_key.public_key()
+
+    @classmethod
+    def from_private_key_file(cls, filename, password=None):
+        with open(filename, "rb") as f:
+            data = f.read()
+        return cls.from_private_key(data, password)
+
+    @classmethod
+    def from_private_key(cls, data, password=None):
+        # Handle file-like object
+        if hasattr(data, "read"):
+            data = data.read()
+
+        if isinstance(data, str):
+            data = data.encode('utf-8')
+
+        p_key = serialization.load_pem_private_key(
+            data,
+            password=password.encode() if isinstance(password, str) else password
+        )
+        if not isinstance(p_key, ed25519.Ed25519PrivateKey):
+             raise ValueError("Not an Ed25519 key")
+
+        return cls(key=p_key)
+
+    def get_name(self) -> str:
+        return "ssh-ed25519"
+
+    def sign(self, data: bytes) -> bytes:
+        if not self.private_key:
+             raise ValueError("Private key not available for signing")
+        return self.private_key.sign(data)
+
+    def verify(self, data: bytes, signature: bytes) -> bool:
+        try:
+            self.public_key.verify(signature, data)
+            return True
+        except InvalidSignature:
+            return False
+
+    def as_bytes(self) -> bytes:
+        """
+        Format: string "ssh-ed25519", string key_bytes
+        """
+        # Get raw bytes
+        pub_bytes = self.public_key.public_bytes(
+            encoding=serialization.Encoding.Raw,
+            format=serialization.PublicFormat.Raw
+        )
+        
+        from ccai.protocols.ssh.sshpacket import SshPacket
+        pkt = SshPacket()
+        pkt.add_string(self.get_name())
+        pkt.add_string(pub_bytes)
+        return pkt.payload

ccai/protocols/ssh/sshpacket.py

@@ -0,0 +1,138 @@
+"""
+SSH Packet Handling
+"""
+
+import struct
+from typing import Union, List
+
+class SshPacket:
+    """
+    Repsrefents an SSH Packet and provides methods for packing/unpacking data types.
+    """
+
+    def __init__(self, message_id: int = None, payload: bytes = b''):
+        self.message_id = message_id
+        self.payload = bytearray(payload)
+        self._position = 0
+
+    @property
+    def remaining(self) -> bytes:
+        return self.payload[self._position:]
+
+    def get_bytes(self, length: int) -> bytes:
+        if self._position + length > len(self.payload):
+            raise IndexError("Not enough data in packet")
+        data = self.payload[self._position : self._position + length]
+        self._position += length
+        return bytes(data)
+
+    def get_byte(self) -> int:
+        return self.get_bytes(1)[0]
+
+    def get_boolean(self) -> bool:
+        return bool(self.get_byte())
+
+    def get_uint32(self) -> int:
+        return struct.unpack('>I', self.get_bytes(4))[0]
+
+    def get_uint64(self) -> int:
+        return struct.unpack('>Q', self.get_bytes(8))[0]
+
+    def get_string(self) -> bytes:
+        """
+        SSH string: uint32 length + content
+        """
+        length = self.get_uint32()
+        return self.get_bytes(length)
+    
+    def get_mpint(self) -> int:
+        """
+        SSH mpint: Two's complement integer, stored as a string.
+        """
+        data = self.get_string()
+        if not data:
+            return 0
+        return int.from_bytes(data, byteorder='big', signed=True)
+
+    def get_name_list(self) -> List[str]:
+        """
+        SSH name-list: comma-separated list of strings
+        """
+        data = self.get_string().decode('utf-8')
+        if not data:
+            return []
+        return data.split(',')
+    
+    # --- Packing Methods ---
+
+    def add_byte(self, b: int):
+        self.payload.append(b & 0xFF)
+
+    def add_bytes(self, data: bytes):
+        """
+        Append raw bytes with no length prefix (e.g. the KEXINIT cookie).
+        """
+        self.payload.extend(data)
+
+    def add_boolean(self, b: bool):
+        self.add_byte(1 if b else 0)
+
+    def add_uint32(self, n: int):
+        self.payload.extend(struct.pack('>I', n))
+
+    def add_uint64(self, n: int):
+        self.payload.extend(struct.pack('>Q', n))
+
+    def add_string(self, s: Union[str, bytes]):
+        if isinstance(s, str):
+            s = s.encode('utf-8')
+        self.add_uint32(len(s))
+        self.payload.extend(s)
+
+    def add_mpint(self, n: int):
+        if n == 0:
+            self.add_uint32(0)
+            return
+        
+        # Calculate number of bytes needed
+        # bit_length() excludes sign, so we might need an extra bit/byte
+        length = (n.bit_length() + 7) // 8
+        try:
+            b = n.to_bytes(length, byteorder='big', signed=True)
+        except OverflowError:
+            length += 1
+            b = n.to_bytes(length, byteorder='big', signed=True)
+            
+        self.add_string(b) # mpint is stored as a string
+
+    def add_name_list(self, names: List[str]):
+        self.add_string(",".join(names))
+
+    def to_bytes(self) -> bytes:
+        """
+        Returns the full payload with message ID if set.
+        Usually the message ID is the first byte of the payload.
+        """
+        if self.message_id is not None:
+             # Prepend message ID if not already in payload? 
+             # Standard practice: explicit message ID + payloads.
+             # but often we build the packet linearly.
+             # If message_id is set in constructor, we assume it's the first byte.
+             return struct.pack('B', self.message_id) + self.payload
+        return bytes(self.payload)
+
+    @classmethod
+    def create_disconnect(cls, reason_code: int, description: str):
+        from ccai.protocols.ssh import sshconstants as constants
+        pkt = cls(constants.SSH_MSG_DISCONNECT)
+        pkt.add_uint32(reason_code)
+        pkt.add_string(description)
+        pkt.add_string("") # Language tag
+        return pkt
+
+    @classmethod
+    def create_ignore(cls, data: bytes):
+        from ccai.protocols.ssh import sshconstants as constants
+        pkt = cls(constants.SSH_MSG_IGNORE)
+        pkt.add_string(data)
+        return pkt