From 2a0ad37a83bfa81231bae88d0ec9b36557de59bf Mon Sep 17 00:00:00 2001
From: Andrew Kay <git@ajk.me>
Date: Mon, 27 Jul 2020 19:01:18 -0500
Subject: [PATCH] Initial support for extended data stream, enough to log into
 TK4

---
 README.md        |  3 +--
 oec/__main__.py  |  6 +++++-
 oec/tn3270.py    | 23 +++++++++++++++++------
 requirements.txt |  2 +-
 4 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 38a49f5..c4184b1 100644
--- a/README.md
+++ b/README.md
@@ -24,10 +24,9 @@ emulation.
 
   - [x] TN3270
       - [x] Basic TN3270
+      - [ ] Extended Data Stream
       - [ ] TN3270E
-      - [ ] EAB (Extended Attribute Buffer)
       - [ ] SSL/TLS
-      - [ ] Non-English character sets
   - [x] VT100
   - [ ] Connection menu
   - [ ] MLT (Multiple Logical Terminals)
diff --git a/oec/__main__.py b/oec/__main__.py
index c23b54f..715b539 100644
--- a/oec/__main__.py
+++ b/oec/__main__.py
@@ -52,7 +52,7 @@ def _create_interface(args):
 
 def _create_session(args, terminal):
     if args.emulator == 'tn3270':
-        return TN3270Session(terminal, args.host, args.port)
+        return TN3270Session(terminal, args.host, args.port, args.extended_data_stream)
 
     if args.emulator == 'vt100' and is_vt100_available:
         host_command = [args.command, *args.command_args]
@@ -91,6 +91,10 @@ def main():
     tn3270_parser.add_argument('host', help='Hostname')
     tn3270_parser.add_argument('port', nargs='?', default=23, type=int)
 
+    tn3270_parser.add_argument('--disable-eds', action='store_false',
+                               dest='extended_data_stream',
+                               help='Disable extended data stream support')
+
     if is_vt100_available:
         vt100_parser = subparsers.add_parser('vt100', description='VT100 emulator',
                                              help='VT100 emulator')
diff --git a/oec/tn3270.py b/oec/tn3270.py
index 2d43acf..7b0a7c8 100644
--- a/oec/tn3270.py
+++ b/oec/tn3270.py
@@ -47,12 +47,13 @@ AID_KEY_MAP = {
 class TN3270Session(Session):
     """TN3270 session."""
 
-    def __init__(self, terminal, host, port):
+    def __init__(self, terminal, host, port, extended_data_stream=True):
         self.logger = logging.getLogger(__name__)
 
         self.terminal = terminal
         self.host = host
         self.port = port
+        self.extended_data_stream = extended_data_stream
 
         self.telnet = None
         self.emulator = None
@@ -154,6 +155,9 @@ class TN3270Session(Session):
     def _connect_host(self):
         terminal_type = f'IBM-3278-{self.terminal.terminal_id.model}'
 
+        if self.extended_data_stream:
+            terminal_type += '-E'
+
         self.telnet = Telnet(terminal_type)
 
         self.telnet.open(self.host, self.port)
@@ -170,9 +174,9 @@ class TN3270Session(Session):
             byte = 0x00
 
             if isinstance(cell, AttributeCell):
-                byte = self._map_attribute(cell.attribute)
+                byte = self._map_attribute(cell)
             elif isinstance(cell, CharacterCell):
-                byte = self._map_character(cell.byte)
+                byte = self._map_character(cell)
 
             self.terminal.display.buffered_write(byte, index=address)
 
@@ -196,17 +200,24 @@ class TN3270Session(Session):
         # TODO: eek, is this the correct place to do this?
         self.operator_error = None
 
-    def _map_attribute(self, attribute):
+    def _map_attribute(self, cell):
         # Only map the protected and display bits - ignore numeric, skip and modified.
-        return 0xc0 | (attribute.value & 0x2c)
+        return 0xc0 | (cell.attribute.value & 0x2c)
+
+    def _map_character(self, cell):
+        byte = cell.byte
 
-    def _map_character(self, byte):
         if byte == DUP:
             return encode_ascii_character(ord('*'))
 
         if byte == FM:
             return encode_ascii_character(ord(';'))
 
+        # TODO: Temporary workaround to show empty reverse video fields until EAB
+        # support is added.
+        if byte == 0x40 and cell.formatting is not None and cell.formatting.reverse:
+            return encode_ascii_character(ord('#'))
+
         return encode_ebcdic_character(byte)
 
     def _format_message_area(self):
diff --git a/requirements.txt b/requirements.txt
index e1c7df3..5515926 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,7 +2,7 @@ ptyprocess==0.6.0
 pycoax==0.5.0
 pyserial==3.4
 pyte==0.8.0
-pytn3270==0.5.1
+pytn3270==0.6.1
 sliplib==0.5.0
 sortedcontainers==2.1.0
 wcwidth==0.1.7