diff --git a/lib/main.dart b/lib/main.dart
index d3d7e82..677e073 100644
--- a/lib/main.dart
+++ b/lib/main.dart
@@ -15,6 +15,29 @@ void main() {
   runApp(const MyApp());
 }
 
+Future<OpenSSHEd25519KeyPair> generateKeyPair() async {
+  final algorithm = Ed25519();
+
+  // Generate a key pair
+  final cryptographyKeyPair = await algorithm.newKeyPair();
+
+  final privateKey = await cryptographyKeyPair.extractPrivateKeyBytes();
+  final publicKey = (await cryptographyKeyPair.extractPublicKey()).bytes;
+
+  final dartsshKeyPair = OpenSSHEd25519KeyPair(
+    Uint8List.fromList(publicKey),
+    Uint8List.fromList(privateKey + publicKey),
+    "",
+  );
+
+  return dartsshKeyPair;
+}
+
+String encodePublicKey(OpenSSHEd25519KeyPair keyPair, {String comment = ''}) {
+  final publicKeyBytes = keyPair.toPublicKey().encode();
+  return "ssh-ed25519 ${base64.encode(publicKeyBytes)} $comment";
+}
+
 class MyApp extends StatelessWidget {
   const MyApp({super.key});
 
@@ -64,35 +87,10 @@ class MyHomePage extends StatefulWidget {
   State<MyHomePage> createState() => _MyHomePageState();
 }
 
-String encodeOpenSshEd25519PublicKey(
-  Uint8List publicKeyBytes, {
-  String comment = '',
-}) {
-  Uint8List encodeString(String str) {
-    final strBytes = utf8.encode(str);
-    final lenBytes = ByteData(4)..setUint32(0, strBytes.length);
-    return Uint8List.fromList(lenBytes.buffer.asUint8List() + strBytes);
-  }
-
-  Uint8List encodeBytes(Uint8List bytes) {
-    final lenBytes = ByteData(4)..setUint32(0, bytes.length);
-    return Uint8List.fromList(lenBytes.buffer.asUint8List() + bytes);
-  }
-
-  final type = 'ssh-ed25519';
-  final keyData = Uint8List.fromList(
-    encodeString(type) + encodeBytes(publicKeyBytes),
-  );
-  final keyBase64 = base64.encode(keyData);
-
-  return '$type $keyBase64${comment.isNotEmpty ? ' $comment' : ''}';
-}
-
 class _MyHomePageState extends State<MyHomePage> {
   int _counter = 0;
   int navIndex = 0;
-  List<int> privateKey = [];
-  List<int> publicKey = [];
+  OpenSSHEd25519KeyPair? keyPair;
   String _output = '';
   String key = '';
 
@@ -109,16 +107,10 @@ class _MyHomePageState extends State<MyHomePage> {
     if (navIndex == 0) {
       SSHSocket? socket;
 
-      OpenSSHEd25519KeyPair? keypair;
-
       try {
-        keypair = OpenSSHEd25519KeyPair(
-          Uint8List.fromList(publicKey),
-          Uint8List.fromList(privateKey + publicKey),
-          "",
-        );
-
-        print(keypair.toPem());
+        if (keyPair != null) {
+          print(keyPair!.toPem());
+        }
 
         socket = await SSHSocket.connect(
           '192.168.2.123',
@@ -137,7 +129,7 @@ class _MyHomePageState extends State<MyHomePage> {
           socket,
           username: 'test',
           //onPasswordRequest: () => '123456',
-          identities: List.of(keypair != null ? [keypair] : []),
+          identities: List.of(keyPair != null ? [keyPair!] : []),
         );
 
         final uptime = await client.run('uptime');
@@ -153,16 +145,9 @@ class _MyHomePageState extends State<MyHomePage> {
 
       print("foo");
     } else {
-      final algorithm = Ed25519();
-
       // Generate a key pair
-      final keyPair = await algorithm.newKeyPair();
 
-      privateKey = await keyPair.extractPrivateKeyBytes();
-      publicKey = (await keyPair.extractPublicKey())
-          .bytes; // Extract the public key bytes
-
-      final publicKeyStr = base64.encode(publicKey);
+      keyPair = await generateKeyPair();
 
       setState(() {
         // private to hex string
@@ -170,10 +155,11 @@ class _MyHomePageState extends State<MyHomePage> {
         //     .map((byte) => byte.toRadixString(16).padLeft(2, '0'))
         //     .join('');
 
-        key = encodeOpenSshEd25519PublicKey(
-          Uint8List.fromList(publicKey),
-          comment: 'leinelab-app-key',
-        );
+        if (keyPair != null) {
+          key = encodePublicKey(keyPair!, comment: "leinelab-app-key");
+        } else {
+          key = "No key pair generated.";
+        }
       });
 
       await SystemChannels.platform.invokeMethod<void>(