#!/usr/bin/env scheme-script
;; -*- mode: scheme; coding: utf-8 -*- !#
;; fcdisasm - The Full-Color Disassembler
;; Copyright © 2008, 2009, 2010 Göran Weinholt <goran@weinholt.se>
;;
;; This program is free software: you can redistribute it and/or modify
;; it under the terms of the GNU General Public License as published by
;; the Free Software Foundation, either version 3 of the License, or
;; (at your option) any later version.
;;
;; This program is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;; GNU General Public License for more details.
;;
;; You should have received a copy of the GNU General Public License
;; along with this program.  If not, see <http://www.gnu.org/licenses/>.
#!r6rs

;; This program is an example of how to use (industria disassembler
;; x86) and a novelty: first disassembler to go *full color* for the
;; hex dump!

(import (rnrs)
        (weinholt binfmt elf (0 (>= 0)))
        #;(prefix (weinholt disassembler arm (1 (>= 0))) arm:)
        (prefix (weinholt disassembler m68hc12 (1 (>= 0))) m68hc12:)
        (prefix (weinholt disassembler mips (1 (>= 0))) mips:)
        (prefix (weinholt disassembler x86 (1 (>= 0))) x86:)
        (only (weinholt disassembler x86 (1 (>= 0))) invalid-opcode?))

;; Print an instruction with hexadecimal numbers.
(define (print-instr/sexpr i)
  (cond ((pair? i)
         (display "(")
         (let lp ((i i))
           (unless (null? i)
             (print-instr/sexpr (car i))
             (unless (null? (cdr i))
               (display #\space)
               (lp (cdr i)))))
         (display ")"))
        ((number? i)
         (display "#x")
         (display (number->string i 16)))
        (else
         (display i))))

(define (disassemble p get-instruction color end-position)
  (let lp ((pos (port-position p)))
    (let* ((tagged-bytes '())
           (i (guard (con
                      ((invalid-opcode? con)
                       (list 'bad:
                             (condition-message con))))
                (get-instruction p
                                 (lambda x
                                   (set! tagged-bytes (cons x tagged-bytes))))))
           (new-pos (port-position p)))
      (unless (or (eof-object? i)
                  (and end-position (> new-pos end-position)))
        (let ((x (number->string pos 16)))
          (if (< (string-length x) 8)
              (display (make-string (- 8 (string-length x)) #\space)))
          (display x))
        (display ": ")
        (for-each (lambda (x)
                    (let ((tag (car x))
                          (bytes (cdr x)))
                      (cond ((eq? tag '/is4)
                             (when color
                               (display "\x1b;[1;34m"))
                             (display (number->string (bitwise-bit-field (car bytes) 4 8) 16))
                             (when color
                               (display "\x1b;[1;37m"))
                             (display (number->string (bitwise-bit-field (car bytes) 0 4) 16)))
                            (else
                             (when color
                               (case tag
                                 ((modr/m sib tfr/exg/sex) (display "\x1b;[1;34m"))
                                 ((opcode) (display "\x1b;[1;32m"))
                                 ((prefix) (display "\x1b;[1;33m"))
                                 ((immediate) (display "\x1b;[1;37m"))
                                 ((disp offset) (display "\x1b;[1;35m"))
                                 (else (display "\x1b;[0m"))))
                             (for-each (lambda (byte)
                                         (when (< byte #x10)
                                           (display #\0))
                                         (display (number->string byte 16)))
                                       bytes)))))
                  (reverse tagged-bytes))
        (when color
          (display "\x1b;[0m"))
        (display (make-string (- 31 (* 2 (- new-pos pos))) #\space))
        (print-instr/sexpr i)
        (newline)
        (lp new-pos)))))

(define (get-disassembler machine endianness)
  (cond #;((= machine EM-ARM) (lambda (p c) (arm:get-instruction p #f c)))
        ((= machine EM-386) (lambda (p c) (x86:get-instruction p 32 c)))
        ((= machine EM-X86-64) (lambda (p c) (x86:get-instruction p 64 c)))
        ((= machine EM-68HC12) (lambda (p c) (m68hc12:get-instruction p c)))
        ((= machine EM-MIPS) (lambda (p c) (mips:get-instruction p endianness c)))
        (else
         (display "No support for this architecture: ")
         (cond ((assv machine EM-names) =>
                (lambda (n) (display (cdr n))))
               (else
                (display "Unknown architecture (")
                (display machine)
                (display ")")))
         (newline)
         (exit))))

(define (disassemble-file filename mode color)
  (cond ((is-elf-image? filename)
         (display "ELF image detected. Looking for .text section...\n")
         (let ((image (open-elf-image-file filename)))
           (cond ((elf-input-file-section-by-name image ".text") =>
                  (lambda (text)
                    (let ((get-instruction (get-disassembler
                                            (elf-input-file-machine image)
                                            (if (= (elf-input-file-endianness image) 1)
                                                (endianness little) (endianness big)))))
                      (set-port-position! (elf-input-file-port image)
                                          (elf-section-header-offset text))
                      (disassemble (elf-input-file-port image)
                                   get-instruction color
                                   (+ (elf-section-header-offset text)
                                      (elf-section-header-size text))))))
                 (else
                  (display "This ELF image has no .text section. No disassembly for you.\n")))))
        (else
         (disassemble (open-file-input-port filename)
                      (lambda (p c) (x86:get-instruction p (or mode 16) c))
                      color #f))))

(define (parse-args args)
  (define (help . msg)
    (let ((x (current-error-port)))
      (when msg (display (car msg) x) (newline x) (newline x))
      (display "fcdisasm - Full-color disassembler

Usage: fcdisasm [-b|--bits <bits>] [--nocolor] [--] <filename>

The <bits> argument can be either 16, 32 or 64 (default).
The --nocolor flag suppresses the color output.

The colors used are:
 * Blue for ModR/M and SIB bytes
 * Green for opcode bytes
 * Yellow for prefix bytes
 * White for immediates
 * Magenta for offsets.

Author: Göran Weinholt <goran@weinholt.se>.
" x)
      (exit 1)))
  (let lp ((filename #f)
           (color #t)
           (mode #f)
           (args args))
    (cond ((null? args)
           (unless filename
             (help "ERROR: No filename given."))
           (values filename mode color))
          ((or (string=? (car args) "--bits")
               (string=? (car args) "-b"))
           (if (null? (cdr args)) (help "ERROR: -b needs an argument (16, 32, 64)"))
           (cond ((assoc (cadr args) '(("64" . 64) ("32" . 32) ("16" . 16))) =>
                  (lambda (x)
                    (lp filename color (cdr x) (cddr args))))
                 (else
                  (help "ERROR: invalid argument for -b flag"))))
          ((string=? (car args) "--nocolor")
           (lp filename #f mode (cdr args)))
          ((string=? (car args) "--")
           (if (not (= (length args) 2)) (help "ERROR: following -- must be only a filename"))
           (if filename (help "ERROR: you can't have it both ways, use -- or don't"))
           (lp (cadr args) color mode (cddr args)))
          (else
           (if filename (help "ERROR: extra arguments on command line"))
           (lp (car args) color mode (cdr args))))))

(define (main args)
  (call-with-values (lambda () (parse-args args))
    disassemble-file))

(main (cdr (command-line)))
