local nmap      = require "nmap"
local stdnse    = require "stdnse"
local shortport = require "shortport"
local tn3270    = require "tn3270"
local brute     = require "brute"
local creds     = require "creds"
local unpwdb    = require "unpwdb"
local io        = require "io"
local table     = require "table"
local string   = require "string"
local stringaux = require "stringaux"


description = [[
CICS transaction ID enumerator for IBM mainframes.
This script is based on mainframe_brute by Dominic White
(https://github.com/sensepost/mainframe_brute). However, this script
doesn't rely on any third party libraries or tools and instead uses
the NSE TN3270 library which emulates a TN3270 screen in lua.

CICS only allows for 4 byte transaction IDs, that is the only specific rule
found for CICS transaction IDs.
]]

---
-- @args idlist Path to list of transaction IDs.
--  Defaults to the list of CICS transactions from IBM.
-- @args cics-enum.commands Commands in a semi-colon separated list needed
--  to access CICS. Defaults to <code>CICS</code>.
-- @args cics-enum.path Folder used to store valid transaction id 'screenshots'
--  Defaults to <code>None</code> and doesn't store anything.
-- @args cics-enum.user Username to use for authenticated enumeration
-- @args cics-enum.pass Password to use for authenticated enumeration
--
-- @usage
-- nmap --script=cics-enum -p 23 <targets>
--
-- nmap --script=cics-enum --script-args=idlist=default_cics.txt,
-- cics-enum.command="exit;logon applid(cics42)",
-- cics-enum.path="/home/dade/screenshots/",cics-enum.noSSL=true -p 23 <targets>
--
-- @output
-- PORT   STATE SERVICE
-- 23/tcp open  tn3270
-- | cics-enum:
-- |   Accounts:
-- |     CBAM: Valid - CICS Transaction ID
-- |     CETR: Valid - CICS Transaction ID
-- |     CEST: Valid - CICS Transaction ID
-- |     CMSG: Valid - CICS Transaction ID
-- |     CEDA: Valid - CICS Transaction ID
-- |     CEDF: Potentially Valid - CICS Transaction ID
-- |     DSNC: Valid - CICS Transaction ID
-- |_  Statistics: Performed 31 guesses in 114 seconds, average tps: 0
--
-- @changelog
-- 2015-07-04 - v0.1 - created by Soldier of Fortran
-- 2015-11-14 - v0.2 - rewrote iterator
-- 2017-01-22 - v0.3 - added authenticated CICS ID enumeration
-- 2019-02-01 - v0.4 - Removed TN3270E support (breaks location)
--
-- @author Philip Young
-- @copyright Same as Nmap--See https://nmap.org/book/man-legal.html
--

author = "Philip Young aka Soldier of Fortran"
license = "Same as Nmap--See https://nmap.org/book/man-legal.html"
categories = {"intrusive", "brute"}
portrule = shortport.port_or_service({23,992}, "tn3270")

--- Saves the Screen generated by the CICS command to disk
--
-- @param filename string containing the name and full path to the file
-- @param data contains the data
-- @return status true on success, false on failure
-- @return err string containing error message if status is false
local function save_screens( filename, data )
  local f = io.open( filename, "w")
  if not f then return false, ("Failed to open file (%s)"):format(filename) end
  if not(f:write(data)) then return false, ("Failed to write file (%s)"):format(filename) end
  f:close()
  return true
end

Driver = {
  new = function(self, host, port, options)
    local o = {}
    setmetatable(o, self)
    self.__index = self
    o.host = host
    o.port = port
    o.options = options
    o.tn3270 = tn3270.Telnet:new()
    o.tn3270:disable_tn3270e()
    return o
  end,
  connect = function( self )
    local status, err = self.tn3270:initiate(self.host,self.port)
    self.tn3270:get_screen_debug(2)
    if not status then
      stdnse.debug("Could not initiate TN3270: %s", err )
      return false
    end
    return true
  end,
  disconnect = function( self )
    self.tn3270:disconnect()
    self.tn3270 = nil
    return true
  end,
  login = function (self, user, pass) -- pass is actually the CICS transaction we want to try
    local commands = self.options['key1']
    local path = self.options['key2']
    local cics_user = self.options['user']
    local cics_pass = self.options['pass']
    local timeout = 300
    local max_blank = 1
    local loop = 1
    local err, status
    stdnse.debug(2,"Getting to CICS")
    local run = stringaux.strsplit(";%s*", commands)
    for i = 1, #run do
      stdnse.debug(1,"Issuing Command (#%s of %s): %s", i, #run ,run[i])
      self.tn3270:send_cursor(run[i])
      self.tn3270:get_all_data()
      self.tn3270:get_screen_debug(2)
    end
    while self.tn3270:isClear() and max_blank < 7 do
      stdnse.debug(2, "Screen is not clear for %s. Reading all data with a timeout of %s. Count %s",pass, timeout, max_blank)
      self.tn3270:get_all_data(timeout)
      timeout = timeout + 100
      max_blank = max_blank + 1
    end

    while not self.tn3270:isClear() and loop < 10 do
      -- by this point we're at *some* CICS transaction
      -- so we send F3 to exit it
      stdnse.debug(2,"Sending: F3")
      self.tn3270:send_pf(3) -- send F3
      self.tn3270:get_all_data()
      self.tn3270:get_screen_debug(2)
      -- now we want to clear the screen
      self.tn3270:send_clear()
      self.tn3270:get_all_data()
      stdnse.debug(2,"Current CLEARed Screen. Loop: %s", loop )
      self.tn3270:get_screen_debug(2)
      loop = loop + 1
    end

    if loop == 10 then
      -- something is wrong but we can still try transactions. Print error to debug.
      stdnse.debug('Error. Failed to get to a blank screen under CICS (sending F3 followed by CLEAR). Try lowering maxthreads to fix.')
    end
    -- If username/password provided try to authenticate first
    if not (cics_user == nil and cics_pass == nil) then -- We're doing authenticated CICS testing now baby!
      stdnse.debug(2,'Logging in with %s / %s for auth testing', cics_user, cics_pass)
      self.tn3270:send_cursor('CESN')
      self.tn3270:get_all_data()
      self.tn3270:get_screen_debug(2)
      local fields = self.tn3270:writeable() -- Get the writeable field areas
      local user_loc = {fields[1][1],cics_user}   -- This is the 'UserID:' field
      local pass_loc = {fields[3][1],cics_pass}   -- This is the 'Password:' field ([2] is a group ID)
      stdnse.debug(2,'Trying CICS: %s : %s', user, pass)
      self.tn3270:send_locations({user_loc,pass_loc})
      self.tn3270:get_all_data()
      stdnse.debug(2,"Screen Received for User ID: %s / %s", user, pass)
      self.tn3270:get_screen_debug(2)
      local count = 1
      while not self.tn3270:find('DFHCE3549') and count < 6 do -- some systems show a message for a bit before we get to CICS again
          self.tn3270:get_all_data(1000) -- loop for 6 seconds
          count = count + 1
      end
    end
    self.tn3270:get_screen_debug(2)
    self.tn3270:send_clear()
    self.tn3270:get_all_data()
    self.tn3270:get_screen_debug(2)
    stdnse.verbose("Trying Transaction ID: %s", pass)
    self.tn3270:send_cursor(pass)
    self.tn3270:get_all_data()

    max_blank = 1
    while self.tn3270:isClear() and max_blank < 7 do
      stdnse.debug(2, "Screen is not clear for %s. Reading all data with a timeout of %s. Count %s",pass, timeout, max_blank)
      self.tn3270:get_all_data(timeout)
      timeout = timeout + 100
      max_blank = max_blank + 1
    end

    stdnse.debug(2,"Screen Received for Transaction ID: %s", pass)
    self.tn3270:get_screen_debug(2)
    if self.tn3270:find('not recognized') or self.tn3270:find('DFHAC2002') then -- known invalid command
      stdnse.debug("Invalid CICS Transaction ID: %s", string.upper(pass))
      return false,  brute.Error:new( "Incorrect CICS Transaction ID" )
    elseif self.tn3270:isClear() then
      stdnse.debug(2,"Empty Screen when we expect an error.")
      -- this can mean that the transaction ID was valid
      -- but it didn't send a screen update so you should check by hand.
      -- We're not dumping this screen to disk because it's blank.
      return true, creds.Account:new("CICS ID [blank screen]", string.upper(pass), creds.State.VALID)
    elseif self.tn3270:find('Unauthorized') or self.tn3270:find('DFHAC2002') then
      -- this is a VALID cics transaction but you must be authenticated to used it
      -- This will be the same screen for each so we dont bother saving it either
      stdnse.verbose("Valid CICS Transaction ID [requires auth]: %s", string.upper(pass))
      return true, creds.Account:new("CICS ID [requires auth]", string.upper(pass), creds.State.VALID)
    elseif self.tn3270:find('DFHAC2008') or self.tn3270:find('DFHAC2206') or self.tn3270:find('DFHAC2028') or
           self.tn3270:find('DFHRT4415') or self.tn3270:find('DFHRT4480') or self.tn3270:find('TSS7254E') then
      -- these are technically valid CICS transactions
      -- but they are of little/no value. If verbosity is turned way up we'll return these/save a screenshot
      -- otherwise there's no point
      -- DFHAC2008 -- TranID has been Disabled
      -- DFHAC2206 -- Abend
      -- DFHRT4415 -- Cannot access through terminal
      -- DFHRT4480 -- No Longer Supported
      -- DFHAC2028 -- cannot be used
      -- TSS7254E  -- Access not available through this facility
      stdnse.verbose("Valid CICS Transaction ID [Abbend or ID Disabled]: %s", string.upper(pass))
      if nmap.verbosity() > 3 then
        if path ~= nil then
          stdnse.verbose(2,"Writting screen to: %s", path..string.upper(pass)..".txt")
          status, err = save_screens(path..string.upper(pass)..".txt",self.tn3270:get_screen())
          if not status then
            stdnse.verbose(2,"Failed writting screen to: %s", path..string.upper(pass)..".txt")
          end
        end
        return true, creds.Account:new("CICS ID [Abbend]", string.upper(pass), creds.State.VALID)
      else
        return false, brute.Error:new( "Correct Transaction ID - Access Denied" )
      end
    elseif not (cics_user == nil and cics_pass == nil) and
           (self.tn3270:find('TSS7251E') or self.tn3270:find('DFHAC2033')) then
      -- We've logged on but we don't have access to this transaction
      -- TSS7251E  : Access Denied to PROGRAM <X>
      -- DFHAC2033 : You are not authorized to use transaction <X>
      stdnse.verbose("Valid CICS Transaction ID [Access Denied]: %s", string.upper(pass))
      if nmap.verbosity() > 3 then
        return true, creds.Account:new("CICS ID [Access Denied]", string.upper(pass), creds.State.VALID)
      else
        return false, brute.Error:new( "Correct Transaction ID - Access Denied" )
      end
    else
      stdnse.verbose("Valid CICS Transaction ID: %s", string.upper(pass))
      if path ~= nil then
        stdnse.verbose(2,"Writting screen to: %s", path..string.upper(pass)..".txt")
        status, err = save_screens(path..string.upper(pass)..".txt",self.tn3270:get_screen())
        if not status then
          stdnse.verbose(2,"Failed writting screen to: %s", path..string.upper(pass)..".txt")
        end
      end
      return true, creds.Account:new("CICS ID", string.upper(pass), creds.State.VALID)
    end
    return false, brute.Error:new("Something went wrong, we didn't get a proper response")
  end
}

--- Tests the target to see if we can even get to CICS
--
-- @param host host NSE object
-- @param port port NSE object
-- @param user CICS userID
-- @param pass CICS userID password
-- @param commands optional script-args of commands to use to get to CICS
-- @return status true on success, false on failure

local function cics_test( host, port, commands, user, pass )
  stdnse.debug("Checking for CICS")
  local tn = tn3270.Telnet:new()
  tn:disable_tn3270e()
  local status, err = tn:initiate(host,port)
  local msg = 'Unable to get to CICS'
  local cics = false -- initially we're not at CICS
  if not status then
    stdnse.debug("Could not initiate TN3270: %s", err )
    return cics
  end
  tn:get_screen_debug(2) -- prints TN3270 screen to debug
  stdnse.debug("Getting to CICS")
  local run = stringaux.strsplit(";%s*", commands)
  for i = 1, #run do
    stdnse.debug(1,"Issuing Command (#%s of %s): %s", i, #run ,run[i])
    tn:send_cursor(run[i])
    tn:get_all_data()
    tn:get_screen_debug(2)
  end
  tn:get_all_data()
  tn:get_screen_debug(2) -- for debug purposes
  -- we should technically be at CICS. So we send:
  --   * F3 to exit the CICS program
  --   * CLEAR (a tn3270 command) to clear the screen.
  --     (you need to clear before sending a transaction ID)
  --   * a known default CICS transaction ID with predictable outcome
  --     (CESF with 'Sign-off is complete.' as the result)
  -- to confirm that we were in CICS. If so we return true
  -- otherwise we return false
  local count = 1
  while not tn:isClear() and count < 6 do
    -- some systems will just kick you off others are slow in responding
    -- this loop continues to try getting out of CICS 6 times. If it can't
    -- then we probably weren't in CICS to begin with.
    if tn:find("Signon") then
      stdnse.debug(2,"Found 'Signon' sending PF3")
      tn:send_pf(3)
      tn:get_all_data()
    end
    tn:get_all_data()
    stdnse.debug(2,"Clearing the Screen")
    tn:send_clear()
    tn:get_all_data()
    tn:get_screen_debug(2)
    count = count + 1
  end
  if count == 6 then
    return cics
  end
  stdnse.debug(2,"Sending CESF (CICS Default Sign-off)")
  tn:send_cursor('CESF')
  tn:get_all_data()
  if tn:isClear() then
    tn:get_all_data(1000)
  end
  tn:get_screen_debug(2)

  if tn:find('off is complete.') then
      cics = true
  end

  if not (user == nil and pass == nil) then -- We're doing authenticated CICS testing now baby!
    stdnse.verbose(2,'Logging in with %s / %s for auth testing', user, pass)
    tn:send_clear()
    tn:get_all_data()
    tn:get_screen_debug(2)
    tn:send_cursor('CESN')
    tn:get_all_data()
    tn:get_screen_debug(2)
    local fields = tn:writeable() -- Get the writeable field areas
    local user_loc = {fields[1][1],user}   -- This is the 'UserID:' field
    local pass_loc = {fields[3][1],pass}   -- This is the 'Password:' field ([2] is a group ID)
    stdnse.verbose('Trying CICS: %s : %s', user, pass)
    tn:send_locations({user_loc,pass_loc})
    tn:get_all_data()
    stdnse.debug(2,"Screen Received for User ID: %s / %s", user, pass)
    tn:get_screen_debug(2)
    count = 1
    while not tn:find('DFHCE3549') and count < 6 do
	      tn:get_all_data(1000) -- loop for 6 seconds
        tn:get_screen_debug(2)
        count = count + 1
    end
    if not tn:find('DFHCE3549') then
        cics = false
        msg = 'Unable to access CICS with User: '..user..' / Pass: '..pass
    else
        tn:send_cursor('CESF')
        tn:get_all_data()
    end
  end

  tn:disconnect()
  return cics,msg
end

-- Filter iterator for unpwdb
-- CICS is limited to 4 characters.
local valid_cics = function(x)
  return (string.len(x) <= 4)
end

function iter(t)
  local i, val
  return function()
    i, val = next(t, i)
    return val
  end
end

action = function(host, port)
  local cics_id_file = stdnse.get_script_args("idlist")
  local path = stdnse.get_script_args(SCRIPT_NAME .. '.path') -- Folder for screenshots
  local commands = stdnse.get_script_args(SCRIPT_NAME .. '.commands') or 'cics'-- VTAM commands/macros to get to CICS
  local username = stdnse.get_script_args(SCRIPT_NAME .. '.user') or nil
  local password = stdnse.get_script_args(SCRIPT_NAME .. '.pass') or nil
  local cics_ids = {"CADP", "CATA", "CATD", "CATR", "CBAM", "CCIN", "CCRL", "CDBC", "CDBD",
    "CDBF", "CDBI", "CDBM", "CDBN", "CDBO", "CDBQ", "CDBT", "CDFS", "CDST",
    "CDTS", "CEBR", "CEBT", "CECI", "CECS", "CEDA", "CEDB", "CEDC", "CEDF",
    "CEDX", "CEGN", "CEHP", "CEHS", "CEKL", "CEMN", "CEMT", "CEOT", "CEPD",
    "CEPF", "CEPH", "CEPM", "CEPQ", "CEPS", "CEPT", "CESC", "CESD", "CESF",
    "CESL", "CESN", "CEST", "CETR", "CEX2", "CFCL", "CFCR", "CFOR", "CFQR",
    "CFQS", "CFTL", "CFTS", "CGRP", "CHLP", "CIDP", "CIEP", "CIND", "CIS1",
    "CIS4", "CISB", "CISC", "CISD", "CISE", "CISM", "CISP", "CISQ", "CISR",
    "CISS", "CIST", "CISU", "CISX", "CITS", "CJLR", "CJSA", "CJSL", "CJSR",
    "CJTR", "CKAM", "CKBC", "CKBM", "CKBP", "CKBR", "CKCN", "CKDL", "CKDP",
    "CKQC", "CKRS", "CKRT", "CKSD", "CKSQ", "CKTI", "CLDM", "CLQ2", "CLR1",
    "CLR2", "CLS1", "CLS2", "CLS3", "CLS4", "CMAC", "CMPX", "CMSG", "CMTS",
    "COVR", "CPCT", "CPIA", "CPIH", "CPIL", "CPIQ", "CPIR", "CPIS", "CPLT",
    "CPMI", "CPSS", "CQPI", "CQPO", "CQRY", "CRLR", "CRMD", "CRMF", "CRPA",
    "CRPC", "CRPM", "CRSQ", "CRSR", "CRST", "CRSY", "CRTE", "CRTP", "CRTX",
    "CSAC", "CSCY", "CSFE", "CSFR", "CSFU", "CSGM", "CSHA", "CSHQ", "CSHR",
    "CSKP", "CSMI", "CSM1", "CSM2", "CSM3", "CSM5", "CSNC", "CSNE", "CSOL",
    "CSPG", "CSPK", "CSPP", "CSPQ", "CSPS", "CSQC", "CSRK", "CSRS", "CSSF",
    "CSSY", "CSTE", "CSTP", "CSXM", "CSZI", "CTIN", "CTSD", "CVMI", "CWBA",
    "CWBG", "CWTO", "CWWU", "CWXN", "CWXU", "CW2A", "CXCU", "CXRE", "CXRT",
    "DSNC"} -- Default CICS from https://www-01.ibm.com/support/knowledgecenter/SSGMCP_5.2.0/com.ibm.cics.ts.systemprogramming.doc/topics/dfha726.html

  cics_id_file = ( (cics_id_file and nmap.fetchfile(cics_id_file)) or cics_id_file )

  if cics_id_file then
    for l in io.lines(cics_id_file) do
      if not l:match("#!comment:") then
        table.insert(cics_ids, l)
      end
    end
  end
  local cicstst,msg = cics_test(host, port, commands, username, password)
  if cicstst then
    local title = 'CICS Transaction IDs'
    if not(username == nil and password == nil) then title = 'CICS Transaction IDs for User: '.. username end
    local options = { key1 = commands, key2 = path, user = username, pass = password }
    stdnse.debug("Starting CICS Transaction ID Enumeration")
    if path ~= nil then stdnse.verbose(2,"Saving Screenshots to: %s", path) end
    local engine = brute.Engine:new(Driver, host, port, options)
    engine.options.script_name = SCRIPT_NAME
    engine:setPasswordIterator(unpwdb.filter_iterator(iter(cics_ids), valid_cics))
    engine.options.passonly = true
    engine.options:setTitle(title)
    local status, result = engine:start()
    return result
  else
    return msg
  end
end